diff --git a/README.md b/README.md index d07ec4d4..f979f6a3 100644 --- a/README.md +++ b/README.md @@ -74,7 +74,8 @@ To run the program, you must enter the command `cargo run --bin ` a |----------|-------------| | [Getting Started Guide](doc/getting_started.md) | Step-by-step instructions for new users | | [Configuration Guide](doc/configuration.md) | Detailed configuration options and file structure | -| [Tracing System](doc/libazurekvp.md) | Understanding the telemetry and tracing capabilities | +| [KVP Data Exchange](doc/kvp.md) | Pool-file format and Hyper-V interfaces | +| [Diagnostics Contract](doc/diagnostics.md) | Telemetry fields, encodings and consumer behavior | | [End-to-End Testing](doc/e2e_testing.md) | How to perform comprehensive system testing | | [Library Documentation](libazureinit/README.md) | Documentation for the libazureinit library | diff --git a/doc/diagnostics.md b/doc/diagnostics.md new file mode 100644 index 00000000..02e8ee7f --- /dev/null +++ b/doc/diagnostics.md @@ -0,0 +1,271 @@ +# KVP Diagnostics Contract + +This contract defines provisioning telemetry for host-side consumers, including +EG and LPA: operations, observations and artifacts such as logs. The +[KVP contract](kvp.md) defines the underlying pool files and transport. + +## Record Format + +Native diagnostics are appended to guest pool 1. A diagnostic occupies one or +more records. Each key has exactly eleven pipe-delimited fields: + +```text +DIAG|||||||||| +``` + +`DIAG` selects the current format. Do not parse an unsupported diagnostic schema +using this layout. Agent versions identify the producer, not the schema. + +| Field | Description | Format | +|---|---|---| +| `DIAG` | Native diagnostic format identifier | Literal `DIAG` | +| `agent` | Reporting agent, such as `azure-init/0.1.1` | UTF-8 text, conventionally `name/VERSION` | +| `vm_id` | VM identity | UUID | +| `kind` | Operation endpoint or standalone observation | `start`, `finish`, or `event` | +| `name` | Operation or observation, such as `provision:run` or `dmesg` | UTF-8 text | +| `event_id` | Shared by an operation's start and finish; unique to a standalone event | UTF-8 identifier, typically UUID | +| `timestamp` | Emission time | RFC 3339 timestamp | +| `encoding` | Stored value representation | `none`, `zlib+b64`, or `gz+b64` | +| `result` | Reported outcome, when applicable | `success`, `fail`, or empty | +| `duration` | Elapsed seconds, when measured | Finite nonnegative double-precision numeric text, or empty | +| `chunk_index` | Chunk position; zero for a single record | Unsigned decimal integer | + +All fields except `result` and `duration` are required and nonempty. Key fields +cannot contain `|` or NUL; there is no key-field escaping. + +## Timing and Correlation + +| Kind | Result | Duration | Meaning | +|---|---|---|---| +| `start` | Empty | Empty | An operation began | +| `finish` | Required | Required | An operation ended with a reported outcome and elapsed time | +| `event` | Optional | Optional | A standalone observation; outcome and timing are independent | + +A start and finish share an event ID, agent, VM and operation name. A finish +carries its own elapsed duration to allow caller to accurately measure +the operation of interest without relying on the timestamps of the emitted +diagnostics. + +### Timestamps and Durations + +Timestamps must conform to RFC 3339. The writer emits UTC `Z` form with +second, millisecond (default), microsecond or nanosecond precision. + +Durations are finite, nonnegative IEEE 754 double-precision seconds, in +decimal or exponent notation. Empty means absent; zero is a measured duration. +See [Implementation Notes](#implementation-notes) for reader limits and +rounding. + +### Examples + +An operation and a compressed observation, shown as stored key/value records: + +```json +[ + { + "key": "DIAG|azure-init/0.1.1|3f2504e0-4f89-41d3-9a0c-0305e82c3301|start|provision:run|8f3e9c4a-1b2c-4d5e-9f01-234567890abc|2026-08-31T12:34:56.789Z|none|||0", + "value": "starting" + }, + { + "key": "DIAG|azure-init/0.1.1|3f2504e0-4f89-41d3-9a0c-0305e82c3301|finish|provision:run|8f3e9c4a-1b2c-4d5e-9f01-234567890abc|2026-08-31T12:34:57.101Z|none|success|0.312000|0", + "value": "provisioning succeeded" + }, + { + "key": "DIAG|azure-init/0.1.1|3f2504e0-4f89-41d3-9a0c-0305e82c3301|event|example|9f3e9c4a-1b2c-4d5e-9f01-234567890abc|2026-08-31T12:34:57.102Z|zlib+b64|||0", + "value": "eJwLSS0uUSguKcrMS1cwNDIGACxqBQ4=" + } +] +``` + +## Reassembly and Decoding + +The producer encodes the whole payload before splitting it into records. +Group members need not be adjacent or ordered in the pool. + +1. Select the schema from the first key field and validate its metadata. +2. Group by the full key excluding only `chunk_index`. Event ID alone is not + sufficient: it would combine start and finish records. +3. Sort indices numerically and require a unique, contiguous sequence from zero. +4. Concatenate values in that order, then decode according to `encoding`. + +### Encodings + +| Token | Stored value | Decoded content | +|---|---|---| +| `none` | Plain UTF-8 without NUL | Text | +| `zlib+b64` | Standard padded base64 of one zlib stream (RFC 1950) | Bytes | +| `gz+b64` | Standard padded base64 of one gzip member (RFC 1952) | Bytes | + +For compressed encodings, decode base64 before decompressing. Native base64 has +no whitespace. Zlib uses DEFLATE with a 32 KiB window (`wbits=15`) and no preset +dictionary; gzip uses a basic header without optional fields. The tokens are +not aliases. Compression does not imply that the decoded content is text. + +## Limits + +The complete key must fit in 254 UTF-8 bytes, including separators and the +chunk index. The table accounts for this writer's output with the default name +limit. Example widths use the finish record above, with default precision and +a single chunk; they are illustrative, not measured production averages. + +| Field | Example bytes | Maximum emitted bytes | Basis | +|---|---:|---:|---| +| `DIAG` | 4 | 4 | Fixed token | +| `agent` | 16 | 32 | `azure-init/0.1.1`; producer text limit | +| `vm_id` | 36 | 36 | Hyphenated UUID in the example; writer UUID limit | +| `kind` | 6 | 6 | `start`/`event`: 5; `finish`: 6 | +| `name` | 13 | 64 | `provision:run`; configurable producer limit, default 64 | +| `event_id` | 36 | 36 | Hyphenated UUID in the example; writer UUID limit | +| `timestamp` | 24 | 30 | UTC `Z` output: 20/24/27/30 for seconds/ms/us/ns | +| `encoding` | 4 | 8 | `none`: 4; `gz+b64`: 6; `zlib+b64`: 8 | +| `result` | 7 | 7 | Empty: 0; `fail`: 4; `success`: 7 | +| `duration` | 8 | 30 | `0.312000`; up to 20 whole-second digits, a point and 9 fractional digits | +| `chunk_index` | 1 | 4 | `0` through `1022` | +| Ten pipe separators | 10 | 10 | One byte each | +| **Total** | **165** | **267** | Sum of field widths | +| **Space remaining** | **89** | **13 over limit** | Against the 254-byte limit | + +The start and compressed-event examples use 149 and 147 key bytes respectively. +The field bounds total 258 bytes with both default precisions, so not every +combination fits; oversized keys are rejected before writing. Widths count text +bytes, not the in-memory size of a double. These are writer budgets, not +universal widths for every RFC 3339 timestamp or numeric spelling. + +Encoded values are limited to 1,022 bytes per chunk, with at most 1,023 chunks +per payload. Producer budgets do not impose equivalent read limits on other +producers' records, and encoded size does not bound decompressed size. + +## Error Handling + +Preserve records that cannot be interpreted, without silently repairing them. + +There is no total-chunk-count field. Missing trailing records in a plain-text +payload cannot be detected if the remaining indices are contiguous from zero. +Compressed streams additionally allow completion and checksum validation. + +## Provisioning Reports + +The exact key `PROVISIONING_REPORT` identifies a separate, single-record health +report. Its value is pipe-delimited CSV whose fields contain `key=value`. +Fields containing pipes, quotes or newlines use CSV double-quote escaping; +split each decoded field on its first `=`. Field order is not significant for +reading. + +| Field | Description | Format | Required | +|---|---|---|---| +| `result` | Provisioning outcome | `success` or `error` (not diagnostic `fail`) | All reports | +| `agent` | Reporting agent | UTF-8 text | All reports | +| `vm_id` | VM identity | UTF-8 text, usually a UUID | All reports | +| `pps_type` | Pre-provisioning type | `None`, `PreprovisionedOSDisk`, `Running`, `Savable`, or `Unknown` | All reports | +| `timestamp` | Report time | RFC 3339 timestamp | All reports | +| `reason` | Failure explanation | UTF-8 text | Error reports | +| `documentation_url` | Help link for a failure | URL text | No; optional for error reports | +| Other fields | Supporting data | `key=value` text fields | No | + +Other fields are ordered supporting data, including duplicate supporting-data +keys. Required fields cannot be duplicated. On error reports, `reason` and +`documentation_url` cannot be duplicated either. Empty text values remain +supported. Report identities and timestamp spellings are preserved; diagnostic +UUID validation and timestamp output formatting are not imposed on reports. + +Reports replace the prior provisioning result, are not chunked, and must fit +one value. + +## Cloud-init Compatibility + +Cloud-init uses a separate format in the same pool. Current keys include a VM +identity that older keys omit: + +```text +CLOUD_INIT|||||[|] +CLOUD_INIT||||[|] +``` + +The incarnation is numeric, and the VM and event identities are UUIDs. Values +contain `name`, `type`, `ts` and `msg`; finishes add `result` and `duration`, +and chunks add `msg_i`. The value's name and type must match the key. + +Apply the grouping and index checks above to the complete cloud-init base key, +including incarnation and source type. `msg_i` must match the key index, and +other metadata must agree. Join the still-escaped `msg` fragments before +unescaping once: cloud-init can split a JSON escape sequence across records, +so individual chunks need not be valid standalone JSON. + +| Normalized field | Cloud-init source | +|---|---| +| Agent | Literal `CLOUD_INIT` | +| Kind | Source `start` or `finish`; all other types become `event` | +| Name | Source name | +| VM identity | Key field when present; otherwise absent | +| Event identity | Key event UUID | +| Timestamp | `ts`, RFC 3339 with offsets allowed; convert to UTC | +| Outcome | Finish `SUCCESS` becomes `success`; `FAIL` becomes `fail` | +| Duration | Finish's numeric seconds | +| Encoding | Embedded `msg` envelope, when present | + +Do not turn a finish result such as `WARN` into a success or failure; retain it +as unparsed data. Other source types become events without a normalized outcome +or duration. Negative or overflowing finish durations are invalid. + +Cloud-init's [Azure producer](https://github.com/canonical/cloud-init/blob/main/cloudinit/sources/helpers/azure.py) +compresses artifacts with zlib and base64-encodes the result, but labels it +`gz+b64`. The `msg` string contains a JSON object with `encoding` and `data` +fields. + +Our reader reassembles and unescapes `msg`, then base64-decodes `data`, ignoring +ASCII whitespace. It selects zlib or gzip from the compressed bytes rather than +trusting the label, so cloud-init's zlib output can be read correctly. Native +`gz+b64` records remain gzip-only. + +Other envelope encodings are unsupported. Messages without an encoding envelope +are plain text. + +## Implementation Notes + +The following describe `libazureinit-kvp`, not requirements for a consumer's API. + +### Reader Output + +Each decoded group appears at its first record's position. Failed groups retain +their physical records and positions. Reads do not modify the pool or sort by +timestamp. Decoding errors are attached to preserved records: + +| Condition | Error token | +|---|---| +| Unsupported `DIAG` schema | `unsupported_version` | +| Missing index zero or an index gap | `incomplete_group` | +| Repeated index | `duplicate_chunk` | +| Unsupported encoding or invalid payload | `undecodable` | +| Invalid metadata or report | `malformed` | + +Unrelated keys have no decoding error. Invalid pool framing or UTF-8 fails the +whole read. Grouping uses memory proportional to input and decoded content; +there is no decoded-size cap. + +Parsed JSON represents decoded binary payloads as base64 objects. Those bytes +are already decompressed; do not send them to a decompressor again. + +Timestamps normalize to UTC, discarding digits beyond nanoseconds. Durations +are limited to unsigned 64-bit whole seconds plus nanoseconds. Unsigned decimals +with up to nine fractional digits remain exact; other forms round to the nearest +nanosecond. Out-of-range values fail decoding; large durations may lose precision +in JSON. + +Cloud-init durations round to microseconds; source encoding labels are preserved. + +### Writer Choices + +The writer requires UUID event IDs. The name limit defaults to 64 UTF-8 bytes +and is configurable with `DiagnosticWriter::with_max_name_bytes`; oversized +names are rejected, not truncated. Durations use fixed-point seconds with +microsecond precision by default, independently of timestamp precision; +finer digits are discarded. + +Validation precedes writing, but I/O failure may leave a partial batch. Pool +cleanup is explicit. The complete key, including each chunk index, must fit +the [budget](#limits) regardless of the configured name limit. + +Report writers emit success fields as `result`, `agent`, `pps_type`, `vm_id`, +`timestamp`, then extras. Failure order is `result`, `reason`, `agent`, extras, +`pps_type`, `vm_id`, `timestamp`, then the optional documentation URL. Consumers +must not depend on that order. \ No newline at end of file diff --git a/doc/getting_started.md b/doc/getting_started.md index 941888c1..0292e69f 100644 --- a/doc/getting_started.md +++ b/doc/getting_started.md @@ -134,5 +134,6 @@ Both containers will output all logs they have access to in order to better debu ## Next Steps - Review the [Configuration Guide](configuration.md) for detailed configuration options -- Understand the [Tracing System](libazurekvp.md) for monitoring and debugging +- Read the [KVP contract](kvp.md) for pool storage and Hyper-V transport behavior +- Read the [diagnostics contract](diagnostics.md) for telemetry formats and consumer behavior - Explore [End-to-End Testing](e2e_testing.md) for comprehensive testing diff --git a/doc/kvp.md b/doc/kvp.md index 26ef7f9b..090c8b16 100644 --- a/doc/kvp.md +++ b/doc/kvp.md @@ -130,7 +130,7 @@ full details. | hv_kvp_daemon | Upsert + full rewrite | `fcntl` | 512 B (field width) | 2,048 B (field width) | N/A | None | Not checked | Yes (shift + rewrite) | Yes (`kvp_update_mem_state`) | 0–3 | | cloud-init | Append-only | `flock()` | 512 B (field width) | 1,024 B (1,023 + null-terminator) | Truncates | Truncate if `mtime` < boot | Yes | No | No | Pool 1 only | | azure-init (current) | Append-only, batched | `flock()` (via `fs2`) | 512 B (field width) | 1,022 B/chunk | Splits across records | Truncate if `mtime` < boot (no lock) | Zero-padded (implicit) | No | No | Pool 1 only (hardcoded) | -| libazureinit-kvp (planned) | Upsert | `flock()` + `fcntl` | Error if > 254 B | Error if > 1,022 B | Error | Option to truncate if `mtime` < boot (with lock) | Explicit null-terminator | Planned | N/A (direct file I/O) | Any pool (configurable) | +| libazureinit-kvp | Append / upsert / replace | `flock()` + `fcntl` | Error if > 254 B | Error if > 1,022 B | Error | Option to truncate if `mtime` < boot (with lock) | Explicit null-terminator | Yes | N/A (direct file I/O) | Any pool (configurable) | #### flock vs fcntl @@ -239,3 +239,43 @@ no truncation. | `HV_KVP_EXCHANGE_MAX_RECORDS` | 1,024 | Max records per pool file | | `HV_KVP_SAFE_MAX_UTF8_KEY_SIZE` | 255 | 254 UTF-8 bytes + null-terminator; no kernel truncation on write path | | `HV_KVP_SAFE_MAX_UTF8_VALUE_SIZE` | 1,023 | 1,022 UTF-8 bytes + null-terminator; no kernel truncation on write path | + +--- + +## Implementation Notes + +### Locking and Size Policy + +- The `libazureinit-kvp` store acquires `flock` and open-file-description + `fcntl` locks, in that order. +- Safe writes enforce the conservative byte budgets above. +- Full-width writes allow the physical field sizes, including fields without + a terminator. The reader accepts full-width UTF-8 in that case, but such + writes are not guaranteed to survive host transport. + +### Record Updates + +- Appending retains duplicates without a record-count cap. +- Inserting updates a key and collapses its duplicates. +- Inserting new keys and replacing the pool enforce a limit of 1,024 distinct + keys. This is a library policy, not a universal KVP format constraint. +- Map-style reads use the last stored value; physical reads retain every record. +- Deletion may swap a record with the file's tail, changing record order. + +### Cleanup and Read Errors + +- Stale-data cleanup is explicit and compares modification time with system + boot time under the write lock. +- Invalid framing or invalid UTF-8 content fails a read. +- Padding after a NUL is ignored rather than interpreted as text. + +The [diagnostics contract](diagnostics.md) defines telemetry carried in KVP +records. Rust API usage is covered by the crate's generated documentation. + +## References + +- [Linux kernel UAPI](https://github.com/torvalds/linux/blob/master/include/uapi/linux/hyperv.h) +- [Linux KVP driver](https://github.com/torvalds/linux/blob/master/drivers/hv/hv_kvp.c) +- [Linux KVP daemon](https://github.com/torvalds/linux/blob/master/tools/hv/hv_kvp_daemon.c) +- [Cloud-init KVP reporting](https://github.com/canonical/cloud-init/blob/main/cloudinit/reporting/handlers.py) +- [Microsoft Data Exchange overview](https://learn.microsoft.com/en-us/windows-server/virtualization/hyper-v/integration-services-data-exchange) diff --git a/libazureinit-kvp/Cargo.toml b/libazureinit-kvp/Cargo.toml index d58f862d..108a4ddf 100644 --- a/libazureinit-kvp/Cargo.toml +++ b/libazureinit-kvp/Cargo.toml @@ -9,13 +9,16 @@ license = "MIT" description = "Hyper-V KVP (Key-Value Pair) storage library for azure-init." [dependencies] -chrono = { version = "0.4", default-features = false, features = ["clock", "std"] } +base64 = "0.22" +chrono = { version = "0.4", default-features = false, features = ["clock", "serde", "std"] } clap = { version = "4.5.21", features = ["derive"] } csv = "1" +flate2 = "1.0" 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 } diff --git a/libazureinit-kvp/src/cli.rs b/libazureinit-kvp/src/cli.rs index fa347cc7..0fd2e50f 100644 --- a/libazureinit-kvp/src/cli.rs +++ b/libazureinit-kvp/src/cli.rs @@ -1,17 +1,22 @@ // 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; use std::process::ExitCode; +use base64::{engine::general_purpose::STANDARD, Engine as _}; +use chrono::SecondsFormat; use clap::{Parser, Subcommand, ValueEnum}; use serde_json::json; use crate::{ - write_report, KvpError, KvpPool, KvpPoolStore, PoolMode, - ProvisioningReport, ReportPpsType, + write_report, Diagnostic, DiagnosticPayload, DiagnosticReader, + DiagnosticWriter, DurationPrecision, Entry, Kind, KvpError, KvpPool, + KvpPoolStore, PoolMode, ProvisioningReport, ReportPpsType, + PROVISIONING_REPORT_KEY, }; const EXIT_OK: u8 = 0; @@ -19,11 +24,14 @@ const EXIT_NOT_FOUND: u8 = 1; const EXIT_USAGE_OR_VALIDATION: u8 = 2; const EXIT_IO: u8 = 3; -/// Default reporting agent identifier, derived from this crate's version +/// Default reporting agent identifier, derived from this crate's version. const DEFAULT_AGENT: &str = concat!("libazureinit-kvp/", env!("CARGO_PKG_VERSION")); -/// Entry point for the `libazureinit-kvp` binary. +/// Runs the command-line interface using process arguments and standard I/O. +/// +/// Returns the command's exit status. Library callers should use the store, +/// diagnostic or report APIs directly. pub fn run() -> ExitCode { let cli = Cli::parse(); let stdout = io::stdout(); @@ -55,14 +63,18 @@ struct Cli { #[arg(long, global = true)] dir: Option, - /// Use full wire-format key/value limits instead of the safe profile. + /// Use larger key/value limits that may be truncated in host transport. #[arg(long = "unsafe", global = true)] unsafe_mode: bool, /// Emit machine-readable JSON for commands that produce output. - #[arg(long, global = true)] + #[arg(long, global = true, conflicts_with = "text")] json: bool, + /// Emit human-readable text for commands that produce output. + #[arg(long, global = true, conflicts_with = "json")] + text: bool, + #[command(subcommand)] command: Command, } @@ -88,8 +100,21 @@ impl OutputMode { enum Command { /// Print store metadata. Info, - /// Print every record in insertion order as KEY=VALUE lines. - Dump, + /// Print every record in pool order (JSON by default; --text for KEY=VALUE). + /// + /// With --parse, decode diagnostics and reports, preserving other or + /// invalid records as raw entries. Entries stay in pool order. + Dump { + /// Decode diagnostics and provisioning reports (kept in pool order). + #[arg(long)] + parse: bool, + /// Filter diagnostic names by substring; retain reports and raw entries. + #[arg(long, requires = "parse")] + name: Option, + /// Filter diagnostics by kind; retain reports and raw entries. + #[arg(long, value_enum, requires = "parse")] + kind: Option, + }, /// Print key=last_value entries sorted by key. Entries, /// Print the last value for KEY (exit 1 if missing). @@ -102,6 +127,21 @@ enum Command { key: String, value: String, }, + /// Emit a standalone diagnostic with a fresh UUID and current timestamp. + Emit { + /// Event name, e.g. user:create_user. + #[arg(long)] + name: String, + /// Event payload. + #[arg(long)] + message: String, + /// VM UUID (defaults to the current VM's ID). + #[arg(long)] + vm_id: Option, + /// Reporting agent identifier (at most 32 UTF-8 bytes). + #[arg(long, default_value = DEFAULT_AGENT)] + agent: String, + }, /// Replace the pool from KEY=VALUE lines read from --file or stdin. Load { /// Read records from PATH instead of stdin. @@ -205,7 +245,30 @@ impl From for KvpPool { } } +/// Diagnostic kind accepted by `dump --parse --kind`. +#[derive(ValueEnum, Clone, Copy, Debug)] +enum KindArg { + Start, + Finish, + Event, +} + +impl From for Kind { + fn from(value: KindArg) -> Self { + match value { + KindArg::Start => Self::Start, + KindArg::Finish => Self::Finish, + KindArg::Event => Self::Event, + } + } +} + fn dispatch(cli: Cli, stdout: &mut W) -> Result { + if cli.json && cli.text { + return Err(CliError::Usage( + "--json and --text cannot be used together".to_owned(), + )); + } let pool: KvpPool = cli.pool.into(); let mode = if cli.unsafe_mode { PoolMode::Unsafe @@ -222,7 +285,21 @@ 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, name, kind } => { + let output = if cli.text { + OutputMode::Text + } else { + OutputMode::Json + }; + dump( + &store, + stdout, + parse, + name.as_deref(), + kind.map(Kind::from), + output, + ) + } Command::Entries => entries(&store, stdout, output), Command::Read { key } => read(&store, stdout, &key, output), Command::Write { append, key, value } => { @@ -233,6 +310,12 @@ fn dispatch(cli: Cli, stdout: &mut W) -> Result { } Ok(EXIT_OK) } + Command::Emit { + name, + message, + vm_id, + agent, + } => emit(&store, name, message, vm_id, agent), Command::Load { file } => load(&store, file), Command::AppendMultiple { file } => append_multiple(&store, file), Command::Delete { key } => delete(&store, stdout, &key, output), @@ -315,8 +398,15 @@ fn info( fn dump( store: &KvpPoolStore, stdout: &mut W, + parse: bool, + name: Option<&str>, + kind: Option, output: OutputMode, ) -> Result { + if parse { + return diagnostics_entries(store, stdout, name, kind, output); + } + let records = store.dump()?; match output { OutputMode::Text => { @@ -425,6 +515,110 @@ fn is_stale( Ok(if stale { EXIT_OK } else { EXIT_NOT_FOUND }) } +fn diagnostics_entries( + store: &KvpPoolStore, + stdout: &mut W, + name: Option<&str>, + kind: Option, + output: OutputMode, +) -> Result { + let mut entries = DiagnosticReader::new(store.clone()).entries()?; + + if name.is_some() || kind.is_some() { + entries.retain(|entry| match entry { + Entry::Diagnostic(diagnostic) => { + name.is_none_or(|needle| diagnostic.key().name.contains(needle)) + && kind.is_none_or(|wanted| diagnostic.kind() == wanted) + } + Entry::Report(_) | Entry::Raw(_) => true, + }); + } + + match output { + OutputMode::Text => { + for entry in &entries { + writeln!(stdout, "{}", diagnostic_entry_text(entry))?; + } + } + OutputMode::Json => { + let value = serde_json::to_value(&entries) + .expect("diagnostic entries always serialize to JSON"); + writeln_json(stdout, &value)?; + } + } + Ok(EXIT_OK) +} + +fn diagnostic_entry_text(entry: &Entry) -> String { + match entry { + Entry::Diagnostic(diagnostic) => diagnostic_text(diagnostic), + Entry::Report(report) => { + format!("{PROVISIONING_REPORT_KEY}={}", report.encode()) + } + Entry::Raw(raw) => { + let mut line = format!("raw key={} value={}", raw.key, raw.value); + if let Some(error) = raw.error { + let _ = write!(line, " error={error}"); + } + line + } + } +} + +fn diagnostic_text(diagnostic: &Diagnostic) -> String { + let key = diagnostic.key(); + let mut line = + format!("diagnostic kind={} agent={}", diagnostic.kind(), key.agent); + if let Some(vm_id) = &key.vm_id { + let _ = write!(line, " vm_id={vm_id}"); + } + let _ = write!(line, " name={} event_id={}", key.name, key.event_id); + let timestamp = key.timestamp.to_rfc3339_opts(SecondsFormat::AutoSi, true); + let encoding = key + .encoding + .as_ref() + .map(ToString::to_string) + .unwrap_or_else(|| "none".to_owned()); + let _ = write!(line, " timestamp={timestamp} encoding={encoding}"); + let (result, duration) = match diagnostic { + Diagnostic::Start(_) => (None, None), + Diagnostic::Finish(finish) => { + (Some(finish.result), Some(finish.duration)) + } + Diagnostic::Event(event) => (event.result, event.duration), + }; + if let Some(result) = result { + let _ = write!(line, " result={result}"); + } + if let Some(duration) = duration { + let seconds = DurationPrecision::Nanos.format(duration); + let seconds = seconds.trim_end_matches('0').trim_end_matches('.'); + let _ = write!(line, " duration={seconds}s"); + } + match diagnostic.payload() { + DiagnosticPayload::Text(text) => { + let _ = write!(line, " payload={text}"); + } + DiagnosticPayload::Bytes(bytes) => { + let _ = write!(line, " payload_b64={}", STANDARD.encode(bytes)); + } + } + line +} + +fn emit( + store: &KvpPoolStore, + name: String, + message: String, + vm_id: Option, + agent: String, +) -> Result { + let vm_id = resolve_vm_id(vm_id)?; + let writer = DiagnosticWriter::new(store.clone(), agent, vm_id)?; + writer.emit_event(&name, message, None, None, None)?; + Ok(EXIT_OK) +} + fn report_success( store: &KvpPoolStore, vm_id: Option, @@ -492,26 +686,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)? { @@ -746,6 +929,7 @@ mod tests { dir: Some(dir.path().to_path_buf()), unsafe_mode: false, json: false, + text: false, command, } } @@ -756,6 +940,7 @@ mod tests { dir: Some(dir.path().to_path_buf()), unsafe_mode: false, json: true, + text: false, command, } } @@ -771,6 +956,15 @@ mod tests { (code, String::from_utf8(out).unwrap()) } + /// A plain `dump` command with no diagnostics parsing. + fn dump_cmd() -> Command { + Command::Dump { + parse: false, + name: None, + kind: 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 +1004,57 @@ 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 { .. })); + } + + #[rstest] + #[case::default(None, false)] + #[case::json(Some("--json"), false)] + #[case::text(Some("--text"), true)] + fn dump_output_mode(#[case] flag: Option<&str>, #[case] text: bool) { + let dir = TempDir::new().unwrap(); + store_at(&dir).append("key", "value").unwrap(); + let mut args = vec![ + "libazureinit-kvp", + "--dir", + dir.path().to_str().unwrap(), + "dump", + ]; + args.extend(flag); + let (code, output) = run_dispatch(Cli::parse_from(args)); + assert_eq!(code, EXIT_OK); + if text { + assert_eq!(output, "key=value\n"); + } else { + assert_eq!( + parse_json(&output), + json!([{"key":"key","value":"value"}]) + ); + } + } + + #[rstest] + #[case::same_scope(vec!["dump", "--json", "--text"])] + #[case::global_json(vec!["--json", "dump", "--text"])] + #[case::global_text(vec!["--text", "dump", "--json"])] + fn json_and_text_are_mutually_exclusive(#[case] args: Vec<&str>) { + match Cli::try_parse_from( + std::iter::once("libazureinit-kvp").chain(args), + ) { + Err(error) => { + assert_eq!( + error.kind(), + clap::error::ErrorKind::ArgumentConflict + ); + } + Ok(cli) => { + let mut output = Vec::new(); + let error = dispatch(cli, &mut output).unwrap_err(); + assert!(matches!(error, CliError::Usage(_))); + assert_eq!(error.exit_code(), EXIT_USAGE_OR_VALIDATION); + assert!(output.is_empty()); + } + } } #[test] @@ -1259,6 +1503,7 @@ mod tests { dir: Some(dir.path().to_path_buf()), unsafe_mode: true, json: false, + text: false, command: Command::Info, }; let (code, out) = run_dispatch(invocation); @@ -1306,8 +1551,14 @@ mod tests { }, )); - let (_, dumped) = run_dispatch(cli(&dir, Command::Dump)); - assert_eq!(dumped, "k=1\nk=2\n"); + let (_, dumped) = run_dispatch(cli(&dir, dump_cmd())); + assert_eq!( + parse_json(&dumped), + json!([ + {"key": "k", "value": "1"}, + {"key": "k", "value": "2"}, + ]) + ); } #[test] @@ -1330,12 +1581,16 @@ mod tests { store.insert("b", "two").unwrap(); store.insert("a", "one").unwrap(); - let (_, dumped) = run_dispatch(cli(&dir, Command::Dump)); - assert!(dumped.contains("a=one")); - assert!(dumped.contains("b=two")); + let (_, dumped) = run_dispatch(cli(&dir, dump_cmd())); + assert_eq!( + parse_json(&dumped), + json!([ + {"key": "b", "value": "two"}, + {"key": "a", "value": "one"}, + ]) + ); let (_, entries) = run_dispatch(cli(&dir, Command::Entries)); - // entries are sorted by key assert_eq!(entries, "a=one\nb=two\n"); } @@ -1370,8 +1625,15 @@ mod tests { .unwrap(); assert_eq!(code, EXIT_OK); - let (_, dumped) = run_dispatch(cli(&dir, Command::Dump)); - assert_eq!(dumped, "a=1\na=2\nb=3\n"); + let (_, dumped) = run_dispatch(cli(&dir, dump_cmd())); + assert_eq!( + parse_json(&dumped), + json!([ + {"key": "a", "value": "1"}, + {"key": "a", "value": "2"}, + {"key": "b", "value": "3"}, + ]) + ); } #[test] @@ -1395,8 +1657,8 @@ mod tests { assert_eq!(code, EXIT_OK); assert_eq!(out, "2\n"); - let (_, dumped) = run_dispatch(cli(&dir, Command::Dump)); - assert_eq!(dumped, "b=2\n"); + let (_, dumped) = run_dispatch(cli(&dir, dump_cmd())); + assert_eq!(parse_json(&dumped), json!([{"key": "b", "value": "2"}])); } #[rstest] @@ -1478,6 +1740,7 @@ mod tests { dir: Some(blocker), unsafe_mode: false, json: false, + text: false, command: Command::Info, }; let mut out = Vec::new(); @@ -1563,6 +1826,17 @@ mod tests { assert_eq!(io_err.exit_code(), EXIT_IO); } + #[rstest] + #[case(KvpError::EmptyEventField { field: "name" })] + #[case(KvpError::EventFieldTooLong { field: "agent", max: 32, actual: 33 })] + #[case(KvpError::InvalidUuid { field: "event_id" })] + #[case(KvpError::TooManyChunks { max: 1023 })] + #[case(KvpError::PayloadNotUtf8)] + #[case(KvpError::UnsupportedEncoding { token: "zstd+b64".into() })] + fn diagnostic_errors_use_validation_exit_code(#[case] error: KvpError) { + assert_eq!(CliError::from(error).exit_code(), EXIT_USAGE_OR_VALIDATION); + } + #[test] fn cli_error_from_conversions() { let from_kvp: CliError = KvpError::EmptyKey.into(); @@ -1602,7 +1876,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); @@ -1614,6 +1888,109 @@ mod tests { assert_eq!(array[2]["value"], "one"); } + #[test] + fn dispatch_parsed_dump_filters_only_diagnostics() { + let dir = TempDir::new().unwrap(); + let store = store_at(&dir); + store.append("note", "raw value").unwrap(); + let writer = DiagnosticWriter::new( + store.clone(), + "agent", + "00000000-0000-0000-0000-000000000abc", + ) + .unwrap(); + writer + .emit_event("skip", "hidden", None, None, None) + .unwrap(); + writer + .emit_event("keep", "visible", None, None, None) + .unwrap(); + let report = + ProvisioningReport::success("agent", "vm-id", ReportPpsType::None); + write_report(&store, &report).unwrap(); + store.append("DIAG_V2|future", "preserved").unwrap(); + + let (_, output) = run_dispatch(cli( + &dir, + Command::Dump { + parse: true, + name: Some("keep".into()), + kind: None, + }, + )); + let entries = parse_json(&output); + let entries = entries.as_array().unwrap(); + assert_eq!(entries.len(), 4); + assert_eq!( + entries[0], + json!({"type": "raw", "key": "note", "value": "raw value"}) + ); + assert_eq!(entries[1]["type"], "diagnostic"); + assert_eq!(entries[1]["name"], "keep"); + assert_eq!(entries[1]["payload"], "visible"); + assert_eq!( + entries[2], + serde_json::to_value(Entry::Report(report)).unwrap() + ); + assert_eq!( + entries[3], + json!({ + "type": "raw", "key": "DIAG_V2|future", "value": "preserved", + "error": "unsupported_version", + }) + ); + } + + #[test] + fn dispatch_parsed_dump_filters_by_kind() { + let dir = TempDir::new().unwrap(); + let store = store_at(&dir); + let event_id = "8f3e9c4a-1b2c-4d5e-9f01-234567890abc"; + let ts = "2026-08-31T12:34:56.789Z"; + let vm = "00000000-0000-0000-0000-000000000abc"; + let diag = |kind: &str| { + format!("DIAG|agent|{vm}|{kind}|span|{event_id}|{ts}|none|||0") + }; + store.append(&diag("start"), "starting").unwrap(); + store.append(&diag("event"), "obs").unwrap(); + store.append("note", "raw").unwrap(); + + let (_, output) = run_dispatch(cli( + &dir, + Command::Dump { + parse: true, + name: None, + kind: Some(KindArg::Start), + }, + )); + let entries = parse_json(&output); + let entries = entries.as_array().unwrap(); + assert_eq!(entries.len(), 2); + assert_eq!(entries[0]["kind"], "start"); + assert_eq!(entries[0]["name"], "span"); + assert_eq!( + entries[1], + json!({"type": "raw", "key": "note", "value": "raw"}) + ); + + let (_, output) = run_dispatch(cli( + &dir, + Command::Dump { + parse: true, + name: None, + kind: Some(KindArg::Event), + }, + )); + let entries = parse_json(&output); + let entries = entries.as_array().unwrap(); + assert_eq!(entries.len(), 2); + assert_eq!(entries[0]["kind"], "event"); + assert_eq!( + entries[1], + json!({"type": "raw", "key": "note", "value": "raw"}) + ); + } + #[test] fn dispatch_entries_json_emits_sorted_object() { let dir = TempDir::new().unwrap(); diff --git a/libazureinit-kvp/src/diagnostics/cloud_init.rs b/libazureinit-kvp/src/diagnostics/cloud_init.rs new file mode 100644 index 00000000..c65ccbd7 --- /dev/null +++ b/libazureinit-kvp/src/diagnostics/cloud_init.rs @@ -0,0 +1,738 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +use std::time::Duration; + +use base64::{engine::general_purpose::STANDARD, Engine as _}; +use chrono::{DateTime, Utc}; +use serde_json::{Number, Value}; +use uuid::Uuid; + +use super::diagnostic::{ + DecodeError, Diagnostic, DiagnosticEvent, DiagnosticFinish, DiagnosticKey, + DiagnosticPayload, DiagnosticStart, Encoding, Outcome, +}; +use super::encoding::decompress; + +pub(super) const PREFIX: &str = "CLOUD_INIT"; + +struct CloudInitKey<'a> { + kind: &'a str, + name: &'a str, + vm_id: Option<&'a str>, + event_id: &'a str, +} + +pub(super) fn split_key(key: &str) -> Result<(&str, Option), DecodeError> { + let fields: Vec<_> = key.split('|').collect(); + let indexed = match fields.len() { + 5 => false, + 6 => Uuid::parse_str(fields[5]).is_err(), + 7 => true, + _ => return Err(DecodeError::Malformed), + }; + let (base, index) = if indexed { + let (base, index) = + key.rsplit_once('|').ok_or(DecodeError::Malformed)?; + (base, Some(super::parse_unsigned(index)?)) + } else { + (key, None) + }; + parse_key(base)?; + Ok((base, index)) +} + +fn parse_key(base: &str) -> Result, DecodeError> { + let fields: Vec<_> = base.split('|').collect(); + let (prefix, incarnation, kind, name, vm_id, event_id) = + match fields.as_slice() { + [prefix, incarnation, kind, name, event_id] => { + (*prefix, *incarnation, *kind, *name, None, *event_id) + } + [prefix, incarnation, kind, name, vm_id, event_id] => { + (*prefix, *incarnation, *kind, *name, Some(*vm_id), *event_id) + } + _ => return Err(DecodeError::Malformed), + }; + if prefix != PREFIX + || base.contains('\0') + || fields.iter().any(|field| field.is_empty()) + || !incarnation.bytes().all(|byte| byte.is_ascii_digit()) + { + return Err(DecodeError::Malformed); + } + if let Some(vm_id) = vm_id { + Uuid::parse_str(vm_id).map_err(|_| DecodeError::Malformed)?; + } + Uuid::parse_str(event_id).map_err(|_| DecodeError::Malformed)?; + Ok(CloudInitKey { + kind, + name, + vm_id, + event_id, + }) +} + +pub(super) fn decode_single( + base: &str, + value: &str, +) -> Result { + let key = parse_key(base)?; + let metadata: Value = + serde_json::from_str(value).map_err(|_| DecodeError::Malformed)?; + if metadata.get("msg_i").is_some() { + return Err(DecodeError::Malformed); + } + let message = metadata + .get("msg") + .and_then(Value::as_str) + .ok_or(DecodeError::Malformed)?; + diagnostic(key, &metadata, message.to_owned()) +} + +/// Chunk indices are already ordered and checked by the reader. +pub(super) fn decode_chunks<'a>( + base: &str, + chunks: impl Iterator, +) -> Result { + let key = parse_key(base)?; + let mut metadata = None; + let mut escaped_message = String::new(); + for (index, value) in chunks { + let (mut current, fragment) = chunk_parts(value)?; + if current.get("msg_i").and_then(Value::as_u64) != Some(index) { + return Err(DecodeError::Malformed); + } + current + .as_object_mut() + .ok_or(DecodeError::Malformed)? + .remove("msg_i"); + if metadata.as_ref().is_some_and(|first| first != ¤t) { + return Err(DecodeError::Malformed); + } + metadata.get_or_insert(current); + escaped_message.push_str(fragment); + } + let message = serde_json::from_str(&format!("\"{escaped_message}\"")) + .map_err(|_| DecodeError::Malformed)?; + diagnostic(key, &metadata.ok_or(DecodeError::Malformed)?, message) +} + +fn chunk_parts(value: &str) -> Result<(Value, &str), DecodeError> { + // cloud-init puts `msg` last and slices its escaped JSON string verbatim. + let (prefix, message) = value + .rsplit_once(",\"msg\":") + .ok_or(DecodeError::Malformed)?; + let fragment = message + .strip_prefix('"') + .and_then(|message| message.strip_suffix("\"}")) + .ok_or(DecodeError::Malformed)?; + let metadata = serde_json::from_str(&format!("{prefix}}}")) + .map_err(|_| DecodeError::Malformed)?; + Ok((metadata, fragment)) +} + +fn diagnostic( + source: CloudInitKey<'_>, + metadata: &Value, + message: String, +) -> Result { + if metadata.get("name").and_then(Value::as_str) != Some(source.name) + || metadata.get("type").and_then(Value::as_str) != Some(source.kind) + { + return Err(DecodeError::Malformed); + } + let timestamp = metadata + .get("ts") + .and_then(Value::as_str) + .ok_or(DecodeError::Malformed)?; + let timestamp = DateTime::parse_from_rfc3339(timestamp) + .map_err(|_| DecodeError::Malformed)? + .with_timezone(&Utc); + let (payload, encoding) = + decode_message(message, source.kind == "compressed")?; + let key = DiagnosticKey { + agent: PREFIX.to_owned(), + vm_id: source.vm_id.map(str::to_owned), + name: source.name.to_owned(), + event_id: source.event_id.to_owned(), + timestamp, + encoding, + }; + match source.kind { + "start" => { + if metadata.get("result").is_some() + || metadata.get("duration").is_some() + { + return Err(DecodeError::Malformed); + } + Ok(Diagnostic::Start(DiagnosticStart { key, payload })) + } + "finish" => { + let result = match metadata.get("result").and_then(Value::as_str) { + Some("SUCCESS") => Outcome::Success, + Some("FAIL") => Outcome::Failure, + _ => return Err(DecodeError::Malformed), + }; + let duration = metadata + .get("duration") + .and_then(Value::as_number) + .ok_or(DecodeError::Malformed)?; + Ok(Diagnostic::Finish(DiagnosticFinish { + key, + payload, + result, + duration: duration_us(duration)?, + })) + } + _ => Ok(Diagnostic::Event(DiagnosticEvent { + key, + payload, + result: None, + duration: None, + })), + } +} + +fn duration_us(seconds: &Number) -> Result { + if let Some(seconds) = seconds.as_u64() { + let micros = seconds + .checked_mul(1_000_000) + .ok_or(DecodeError::Malformed)?; + return Ok(Duration::from_micros(micros)); + } + let micros = seconds.as_f64().ok_or(DecodeError::Malformed)? * 1_000_000.0; + // The exclusive upper bound avoids a saturating float-to-integer cast. + if !(0.0..u64::MAX as f64).contains(µs) { + return Err(DecodeError::Malformed); + } + Ok(Duration::from_micros(micros.round() as u64)) +} + +fn decode_message( + message: String, + compressed: bool, +) -> Result<(DiagnosticPayload, Option), DecodeError> { + let envelope = serde_json::from_str::(&message).ok(); + let is_envelope = envelope.as_ref().is_some_and(|value| { + value.get("encoding").is_some() && value.get("data").is_some() + }); + if !compressed && !is_envelope { + return Ok((DiagnosticPayload::Text(message), None)); + } + let envelope = envelope.ok_or(DecodeError::Malformed)?; + let encoding = envelope + .get("encoding") + .and_then(Value::as_str) + .ok_or(DecodeError::Malformed)?; + if encoding != "gz+b64" { + return Err(DecodeError::Undecodable); + } + let data = envelope + .get("data") + .and_then(Value::as_str) + .ok_or(DecodeError::Malformed)?; + Ok((decode_compressed(data)?, Some(Encoding::GzB64))) +} + +fn decode_compressed(data: &str) -> Result { + let compact: Vec<_> = data + .bytes() + .filter(|byte| !byte.is_ascii_whitespace()) + .collect(); + let compressed = STANDARD + .decode(compact) + .map_err(|_| DecodeError::Undecodable)?; + // cloud-init's report_compressed_event uses zlib.compress under `gz+b64`. + let encoding = if compressed.starts_with(&[0x1f, 0x8b]) { + Encoding::GzB64 + } else { + Encoding::ZlibB64 + }; + decompress(&compressed, &encoding) +} + +#[cfg(test)] +mod tests { + use super::*; + use std::fs; + + use rstest::rstest; + use serde_json::json; + use tempfile::TempDir; + + use super::super::diagnostic::{Entry, Kind, RawKeyValue}; + use super::super::encoding::decode_payload; + use super::super::reader::DiagnosticReader; + use crate::{KvpPool, KvpPoolStore, PoolMode}; + + const VM_ID: &str = "0e5e179d-5341-478b-8456-fbb90621bdf8"; + const EVENT_ID: &str = "b7a822ba-4eea-46c0-b559-e84396101132"; + const TIMESTAMP: &str = "2026-07-27T23:33:24.339006+02:00"; + const ZLIB_DATA: &str = "eJxLzskvTdHNzMss4WL4DwAi1AUC"; + const GZIP_DATA: &str = "H4sIAAAAAAAC/0vOyS9N0c3MyyzhYvgPACZ1n10NAAAA"; + + fn key(kind: &str, current: bool, index: Option) -> String { + let identity = if current { + format!("{VM_ID}|{EVENT_ID}") + } else { + EVENT_ID.to_owned() + }; + let base = format!("CLOUD_INIT|100|{kind}|test|{identity}"); + match index { + Some(index) => format!("{base}|{index}"), + None => base, + } + } + + fn value(kind: &str, message: &str) -> Value { + json!({"name": "test", "type": kind, "ts": TIMESTAMP, "msg": message}) + } + + fn chunk(kind: &str, index: u64, fragment: &str) -> String { + format!( + r#"{{"name":"test","type":"{kind}","ts":"{TIMESTAMP}","msg_i":{index},"msg":"{fragment}"}}"# + ) + } + + fn entries(records: &[(String, String)]) -> Vec { + let dir = TempDir::new().unwrap(); + let store = + KvpPoolStore::new_in(KvpPool::Guest, dir.path(), PoolMode::Safe) + .unwrap(); + store.append_multiple(records.iter().cloned()).unwrap(); + let before = fs::read(store.path()).unwrap(); + let entries = DiagnosticReader::new(store.clone()).entries().unwrap(); + assert_eq!(fs::read(store.path()).unwrap(), before); + entries + } + + fn only_diagnostic(mut entries: Vec) -> Diagnostic { + assert_eq!(entries.len(), 1); + let mut diagnostic = None; + if let Some(Entry::Diagnostic(value)) = entries.pop() { + diagnostic = Some(value); + } + diagnostic.expect("expected a diagnostic") + } + + fn assert_raw(records: &[(String, String)], error: DecodeError) { + let expected: Vec<_> = records + .iter() + .map(|(key, value)| { + Entry::Raw(RawKeyValue { + key: key.clone(), + value: value.clone(), + error: Some(error), + }) + }) + .collect(); + assert_eq!(entries(records), expected); + } + + #[rstest] + #[case::old_single(false, None)] + #[case::current_single(true, None)] + #[case::old_chunk(false, Some(0))] + #[case::current_chunk(true, Some(17))] + fn key_layouts_identify_chunk_suffixes( + #[case] current: bool, + #[case] index: Option, + ) { + let key = key("event", current, index); + let (base, parsed_index) = split_key(&key).unwrap(); + let parsed = parse_key(base).unwrap(); + assert_eq!(parsed_index, index); + assert_eq!(parsed.vm_id, current.then_some(VM_ID)); + assert_eq!(parsed.event_id, EVENT_ID); + } + + #[test] + fn malformed_base_key_layout_is_rejected() { + assert!(matches!( + parse_key("CLOUD_INIT|100|event|test"), + Err(DecodeError::Malformed) + )); + } + + #[rstest] + #[case::layout("CLOUD_INIT|100|event".into())] + #[case::vm(key("event", true, None).replace(VM_ID, "invalid"))] + #[case::incarnation(key("event", false, None).replace("|100|", "|bad|"))] + #[case::index(format!("{}|-1", key("event", true, None)))] + fn malformed_cloud_keys_remain_raw(#[case] key: String) { + assert_raw( + &[(key, value("event", "message").to_string())], + DecodeError::Malformed, + ); + } + + #[rstest] + #[case::older(false)] + #[case::current(true)] + fn source_identity_and_timestamp_are_mapped(#[case] current: bool) { + let diagnostic = only_diagnostic(entries(&[( + key("event", current, None), + value("event", "message").to_string(), + )])); + assert_eq!(diagnostic.key().agent, "CLOUD_INIT"); + assert_eq!(diagnostic.key().vm_id.as_deref(), current.then_some(VM_ID)); + assert_eq!(diagnostic.key().name, "test"); + assert_eq!(diagnostic.key().event_id, EVENT_ID); + let rendered = serde_json::to_value(diagnostic).unwrap(); + assert_eq!(rendered["timestamp"], "2026-07-27T21:33:24.339006Z"); + assert!(rendered.get("boot_epoch").is_none()); + assert!(rendered.get("diagnostic_version_id").is_none()); + } + + #[rstest] + #[case::start("start", Kind::Start)] + #[case::event("event", Kind::Event)] + #[case::system_info("system-info", Kind::Event)] + #[case::warning("warning", Kind::Event)] + fn source_type_determines_timeline_kind( + #[case] source: &str, + #[case] kind: Kind, + ) { + let diagnostic = only_diagnostic(entries(&[( + key(source, true, None), + value(source, "observed").to_string(), + )])); + assert_eq!(diagnostic.kind(), kind); + assert_eq!( + diagnostic.payload(), + &DiagnosticPayload::Text("observed".into()) + ); + } + + #[test] + fn finish_closes_its_span_with_mapped_outcome_and_microseconds() { + let mut finish = value("finish", "finished with failure"); + finish["result"] = json!("FAIL"); + finish["duration"] = json!(0.1234); + let entries = entries(&[ + ( + key("start", true, None), + value("start", "starting").to_string(), + ), + (key("finish", true, None), finish.to_string()), + ]); + assert!(matches!( + entries.as_slice(), + [Entry::Diagnostic(Diagnostic::Start(start)), Entry::Diagnostic(Diagnostic::Finish(finish))] + if start.key.event_id == finish.key.event_id + && finish.result == Outcome::Failure + && finish.duration == Duration::from_micros(123_400) + )); + } + + #[test] + fn captured_successful_finish_decodes() { + let records = [( + "CLOUD_INIT|1785187982|finish|modules-final/config-scripts_user|0e5e179d-5341-478b-8456-fbb90621bdf8|e5f01809-a7a3-4279-aa64-1f18e21eda6e".into(), + 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"}"#.into(), + )]; + let diagnostic = only_diagnostic(entries(&records)); + assert!(matches!(&diagnostic, Diagnostic::Finish(finish) + if finish.result == Outcome::Success + && finish.duration == Duration::from_micros(645) + && finish.key.name == "modules-final/config-scripts_user")); + } + + #[test] + fn warn_finish_is_preserved_without_reclassifying_it_as_an_event() { + let mut warning = value("finish", "completed with warnings"); + warning["result"] = json!("WARN"); + warning["duration"] = json!(1); + let raw = RawKeyValue { + key: key("finish", true, None), + value: warning.to_string(), + error: Some(DecodeError::Malformed), + }; + let entries = entries(&[ + ( + key("start", true, None), + value("start", "starting").to_string(), + ), + (raw.key.clone(), raw.value.clone()), + ]); + assert!(matches!( + &entries[0], + Entry::Diagnostic(Diagnostic::Start(_)) + )); + assert_eq!(entries[1], Entry::Raw(raw)); + } + + #[rstest] + #[case::result("result")] + #[case::duration("duration")] + fn finish_requires_result_and_duration(#[case] missing: &str) { + let mut metadata = value("finish", "finished"); + metadata["result"] = json!("SUCCESS"); + metadata["duration"] = json!(1); + metadata.as_object_mut().unwrap().remove(missing); + assert_raw( + &[(key("finish", true, None), metadata.to_string())], + DecodeError::Malformed, + ); + } + + #[rstest] + #[case::invalid_json("event", "not json".into())] + #[case::timestamp("event", value("event", "message").to_string().replace(TIMESTAMP, "bad"))] + #[case::name("event", value("event", "message").to_string().replace("\"test\"", "\"other\""))] + #[case::message("event", json!({"name":"test", "type":"event", "ts":TIMESTAMP, "msg":7}).to_string())] + #[case::start_result("start", { + let mut metadata = value("start", "message"); + metadata["result"] = json!("SUCCESS"); + metadata.to_string() + })] + #[case::start_duration("start", { + let mut metadata = value("start", "message"); + metadata["duration"] = json!(1); + metadata.to_string() + })] + fn malformed_source_values_remain_raw( + #[case] kind: &str, + #[case] value: String, + ) { + assert_raw(&[(key(kind, true, None), value)], DecodeError::Malformed); + } + + #[test] + fn chunk_without_a_key_index_is_malformed() { + assert_raw( + &[(key("event", false, None), chunk("event", 0, "partial"))], + DecodeError::Malformed, + ); + } + + #[rstest] + #[case::newline(r#"line1\"#, "nline2", "line1\nline2")] + #[case::unicode(r"\ud83", r"d\ude00", "😀")] + fn escaped_fragments_are_unescaped_only_after_reassembly( + #[case] first: &str, + #[case] second: &str, + #[case] expected: &str, + ) { + let first = chunk("event", 0, first); + assert!(serde_json::from_str::(&first).is_err()); + let diagnostic = only_diagnostic(entries(&[ + (key("event", true, Some(1)), chunk("event", 1, second)), + (key("event", true, Some(0)), first), + ])); + assert_eq!( + diagnostic.payload(), + &DiagnosticPayload::Text(expected.into()) + ); + } + + #[rstest] + #[case::index(chunk("event", 0, "second"))] + #[case::metadata(chunk("event", 1, "second").replace(TIMESTAMP, "2026-07-27T21:34:00Z"))] + #[case::escape(chunk("event", 1, r"\x"))] + fn invalid_chunk_values_preserve_the_entire_group(#[case] second: String) { + assert_raw( + &[ + (key("event", true, Some(1)), second), + (key("event", true, Some(0)), chunk("event", 0, "first")), + ], + DecodeError::Malformed, + ); + } + + #[rstest] + #[case::gap(2, DecodeError::IncompleteGroup)] + #[case::duplicate(0, DecodeError::DuplicateChunk)] + fn cloud_groups_use_shared_index_validation( + #[case] second: u64, + #[case] error: DecodeError, + ) { + assert_raw( + &[ + (key("event", false, Some(0)), chunk("event", 0, "first")), + ( + key("event", false, Some(second)), + chunk("event", second, "second"), + ), + ], + error, + ); + } + + #[test] + fn incarnation_keeps_otherwise_identical_groups_separate() { + let records = [ + (key("event", false, Some(0)), chunk("event", 0, "a")), + ( + key("event", false, Some(0)).replace("|100|", "|101|"), + chunk("event", 0, "x"), + ), + (key("event", false, Some(1)), chunk("event", 1, "b")), + ( + key("event", false, Some(1)).replace("|100|", "|101|"), + chunk("event", 1, "y"), + ), + ]; + let entries = entries(&records); + assert!(matches!( + entries.as_slice(), + [Entry::Diagnostic(first), Entry::Diagnostic(second)] + if first.key() == second.key() + && first.payload() == &DiagnosticPayload::Text("ab".into()) + && second.payload() == &DiagnosticPayload::Text("xy".into()) + )); + } + + #[test] + fn mixed_sources_keep_first_seen_order_and_unrelated_raw_records() { + let v1 = format!("DIAG|azure-init|{VM_ID}|event|test|{EVENT_ID}|2026-07-27T21:33:00.000Z|none|||0"); + let entries = entries(&[ + (key("event", true, Some(1)), chunk("event", 1, "b")), + ("unrelated".into(), "value".into()), + (v1, "v1 message".into()), + (key("event", true, Some(0)), chunk("event", 0, "a")), + ]); + assert_eq!(entries.len(), 3); + assert!( + matches!(&entries[0], Entry::Diagnostic(d) if d.key().agent == PREFIX) + ); + assert!( + matches!(&entries[1], Entry::Raw(raw) if raw.key == "unrelated" && raw.error.is_none()) + ); + assert!( + matches!(&entries[2], Entry::Diagnostic(d) if d.key().agent == "azure-init") + ); + } + + #[test] + fn cloud_names_are_not_capped_by_writer_budgets() { + let name = "long-subject".repeat(8); + let mut metadata = value("system-info", "message"); + metadata["name"] = json!(name); + let diagnostic = only_diagnostic(entries(&[( + key("system-info", true, None) + .replace("|test|", &format!("|{name}|")), + metadata.to_string(), + )])); + assert_eq!(diagnostic.key().name, name); + } + + #[rstest] + #[case::whole_seconds(json!(2), 2_000_000)] + #[case::zero(json!(0), 0)] + #[case::fraction(json!(0.1234), 123_400)] + #[case::sub_millisecond(json!(0.00064), 640)] + fn duration_conversion_rounds_to_microseconds( + #[case] seconds: Value, + #[case] expected: u64, + ) { + assert_eq!( + duration_us(seconds.as_number().unwrap()).unwrap(), + Duration::from_micros(expected) + ); + } + + #[rstest] + #[case::negative(json!(-1))] + #[case::integer_overflow(json!(u64::MAX))] + #[case::float_overflow(json!(1e30))] + fn duration_conversion_rejects_invalid_ranges(#[case] seconds: Value) { + assert_eq!( + duration_us(seconds.as_number().unwrap()), + Err(DecodeError::Malformed) + ); + } + + #[rstest] + #[case::zlib(ZLIB_DATA)] + #[case::gzip(GZIP_DATA)] + fn compressed_envelopes_decode_arbitrary_bytes(#[case] data: &str) { + let message = json!({"encoding": "gz+b64", "data": data}).to_string(); + let diagnostic = only_diagnostic(entries(&[( + key("compressed", true, None), + value("compressed", &message).to_string(), + )])); + assert_eq!(diagnostic.kind(), Kind::Event); + assert_eq!(diagnostic.key().encoding, Some(Encoding::GzB64)); + assert_eq!( + diagnostic.payload(), + &DiagnosticPayload::Bytes(b"cloud-init\n\x00\xff".to_vec()) + ); + } + + #[test] + fn all_truncated_zlib_prefixes_are_undecodable() { + let compressed = STANDARD.decode(ZLIB_DATA).unwrap(); + for end in 0..compressed.len() { + assert_eq!( + (end, decode_compressed(&STANDARD.encode(&compressed[..end]))), + (end, Err(DecodeError::Undecodable)) + ); + } + } + + #[test] + fn zlib_checksum_and_trailing_data_are_validated() { + let compressed = STANDARD.decode(ZLIB_DATA).unwrap(); + let mut corrupt = compressed.clone(); + *corrupt.last_mut().unwrap() ^= 1; + let mut trailing = compressed; + trailing.push(0); + for bytes in [corrupt, trailing] { + assert_eq!( + decode_compressed(&STANDARD.encode(bytes)), + Err(DecodeError::Undecodable) + ); + } + } + + #[test] + fn compatibility_does_not_relax_native_gzip_validation() { + assert_eq!( + decode_payload(ZLIB_DATA.as_bytes(), Some(&Encoding::GzB64)), + Err(DecodeError::Undecodable) + ); + } + + #[rstest] + #[case::standalone_base64("b64")] + #[case::unknown("zstd")] + fn other_encoding_labels_are_not_supported(#[case] encoding: &str) { + let message = + json!({"encoding": encoding, "data": GZIP_DATA}).to_string(); + assert_raw( + &[( + key("compressed", true, None), + value("compressed", &message).to_string(), + )], + DecodeError::Undecodable, + ); + } + + #[test] + fn compressed_envelope_requires_string_data() { + let message = json!({"encoding": "gz+b64", "data": 7}).to_string(); + assert_raw( + &[( + key("compressed", true, None), + value("compressed", &message).to_string(), + )], + DecodeError::Malformed, + ); + } + + #[test] + fn ordinary_json_messages_are_not_assumed_to_be_encoded() { + let message = r#"{"encoding":"utf-8","description":"ordinary JSON"}"#; + let diagnostic = only_diagnostic(entries(&[( + key("event", false, None), + value("event", message).to_string(), + )])); + assert_eq!( + diagnostic.payload(), + &DiagnosticPayload::Text(message.into()) + ); + assert_eq!(diagnostic.key().encoding, None); + } +} diff --git a/libazureinit-kvp/src/diagnostics/diagnostic.rs b/libazureinit-kvp/src/diagnostics/diagnostic.rs new file mode 100644 index 00000000..c568f891 --- /dev/null +++ b/libazureinit-kvp/src/diagnostics/diagnostic.rs @@ -0,0 +1,690 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +use std::fmt; +use std::time::Duration; + +use base64::{engine::general_purpose::STANDARD, Engine as _}; +use chrono::{DateTime, SecondsFormat, Utc}; +use serde::ser::SerializeStruct; +use serde::{Serialize, Serializer}; + +use crate::ProvisioningReport; + +/// Prefix used for diagnostics emitted by +/// [`DiagnosticWriter`](crate::DiagnosticWriter). +pub const DIAGNOSTIC_VERSION_ID: &str = "DIAG"; + +/// Whether a diagnostic starts or finishes an operation, or records an event. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize)] +#[serde(rename_all = "lowercase")] +pub enum Kind { + /// An operation began. + Start, + /// An operation ended. + Finish, + /// A standalone observation. + Event, +} + +impl fmt::Display for Kind { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str(match self { + Self::Start => "start", + Self::Finish => "finish", + Self::Event => "event", + }) + } +} + +/// Encoding used to compress a diagnostic payload for storage. +/// +/// Pass `None` to [`DiagnosticWriter`](crate::DiagnosticWriter) for plain text. +/// Prefer [`ZlibB64`](Self::ZlibB64) for telemetry consumed by Kusto. +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum Encoding { + /// Zlib followed by standard base64 (`zlib+b64`). + ZlibB64, + /// Gzip followed by standard base64 (`gz+b64`). The reader also accepts + /// cloud-init's zlib data under this label. + GzB64, + /// An unsupported encoding token; writers reject it. + Other(String), +} + +impl fmt::Display for Encoding { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str(match self { + Self::ZlibB64 => "zlib+b64", + Self::GzB64 => "gz+b64", + Self::Other(token) => token, + }) + } +} + +impl Serialize for Encoding { + fn serialize(&self, serializer: S) -> Result + where + S: Serializer, + { + serializer.collect_str(self) + } +} + +/// Reported result of an operation or observation. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize)] +#[serde(rename_all = "lowercase")] +pub enum Outcome { + /// The operation succeeded. + Success, + /// The operation failed, represented as `fail` in serialized output. + #[serde(rename = "fail")] + Failure, +} + +impl fmt::Display for Outcome { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str(match self { + Self::Success => "success", + Self::Failure => "fail", + }) + } +} + +/// The text or bytes carried by a diagnostic. +/// +/// Strings convert to [`Text`](Self::Text); byte slices and vectors convert to +/// [`Bytes`](Self::Bytes). Reading a compressed payload always returns bytes, +/// even when its contents are valid text. +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum DiagnosticPayload { + /// A UTF-8 message. + Text(String), + /// Binary content, serialized to JSON as a base64-encoded object. + Bytes(Vec), +} + +impl From<&str> for DiagnosticPayload { + fn from(value: &str) -> Self { + Self::Text(value.to_owned()) + } +} + +impl From for DiagnosticPayload { + fn from(value: String) -> Self { + Self::Text(value) + } +} + +impl From<&[u8]> for DiagnosticPayload { + fn from(value: &[u8]) -> Self { + Self::Bytes(value.to_vec()) + } +} + +impl From> for DiagnosticPayload { + fn from(value: Vec) -> Self { + Self::Bytes(value) + } +} + +impl Serialize for DiagnosticPayload { + fn serialize(&self, serializer: S) -> Result + where + S: Serializer, + { + match self { + Self::Text(text) => serializer.serialize_str(text), + Self::Bytes(bytes) => { + let mut payload = + serializer.serialize_struct("DiagnosticPayload", 3)?; + payload.serialize_field("type", "bytes")?; + payload.serialize_field("encoding", "base64")?; + payload.serialize_field("data", &STANDARD.encode(bytes))?; + payload.end() + } + } + } +} + +/// Why a recognized diagnostic or report could not be decoded. +/// +/// The reader preserves the original record and attaches this error to +/// [`RawKeyValue::error`]. Failures to read the pool instead return +/// [`KvpError`](crate::KvpError). +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize)] +#[serde(rename_all = "snake_case")] +pub enum DecodeError { + /// The record belongs to an unsupported diagnostic schema. + UnsupportedVersion, + /// Chunk indices are missing or do not start at zero. + IncompleteGroup, + /// More than one record uses the same chunk index. + DuplicateChunk, + /// The payload encoding is unsupported or its contents are invalid. + Undecodable, + /// Diagnostic or report metadata does not match its format. + Malformed, +} + +impl fmt::Display for DecodeError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str(match self { + Self::UnsupportedVersion => "unsupported diagnostic schema version", + Self::IncompleteGroup => { + "diagnostic chunks are not contiguous from index 0" + } + Self::DuplicateChunk => "duplicate diagnostic chunk index", + Self::Undecodable => "diagnostic payload could not be decoded", + Self::Malformed => "malformed diagnostic or provisioning report", + }) + } +} + +impl std::error::Error for DecodeError {} + +/// Producer, identity and timestamp shared by all diagnostic kinds. +#[derive(Clone, Debug, PartialEq, Eq, Serialize)] +pub struct DiagnosticKey { + /// Reporting agent, conventionally `name/VERSION`. + pub agent: String, + /// VM UUID; `None` for older cloud-init records. + pub vm_id: Option, + /// Operation or observation name, such as `provision:run` or `dmesg`. + pub name: String, + /// UTF-8 identifier, typically UUID, shared by an operation's start and finish; + /// unique for a standalone event. + pub event_id: String, + /// When this diagnostic was emitted, in UTC. + #[serde(serialize_with = "serialize_timestamp")] + pub timestamp: DateTime, + /// Stored payload encoding; `None` means UTF-8-encoded plain text. + #[serde(serialize_with = "serialize_encoding")] + pub encoding: Option, +} + +/// The beginning of an operation whose finish uses the same event ID. +#[derive(Clone, Debug, PartialEq, Eq, Serialize)] +pub struct DiagnosticStart { + /// Producer, identity and time of the start. + #[serde(flatten)] + pub key: DiagnosticKey, + /// Message or artifact associated with the start. + pub payload: DiagnosticPayload, +} + +/// The end of an operation, including its outcome and elapsed time. +#[derive(Clone, Debug, PartialEq, Eq, Serialize)] +pub struct DiagnosticFinish { + /// Producer, identity and time of the finish. + #[serde(flatten)] + pub key: DiagnosticKey, + /// Message or artifact associated with the finish. + pub payload: DiagnosticPayload, + /// Reported outcome of the operation. + pub result: Outcome, + /// Elapsed time, serialized to JSON as a number of seconds. + #[serde( + rename = "duration", + serialize_with = "serialize_duration_seconds" + )] + pub duration: Duration, +} + +/// A standalone observation, optionally with an outcome or elapsed time. +#[derive(Clone, Debug, PartialEq, Eq, Serialize)] +pub struct DiagnosticEvent { + /// Producer, identity and time of the observation. + #[serde(flatten)] + pub key: DiagnosticKey, + /// Message or artifact captured by the event. + pub payload: DiagnosticPayload, + /// Reported outcome, when applicable. + #[serde(skip_serializing_if = "Option::is_none")] + pub result: Option, + /// Elapsed time if measured, serialized to JSON as a number of seconds. + #[serde( + rename = "duration", + skip_serializing_if = "Option::is_none", + serialize_with = "serialize_opt_duration_seconds" + )] + pub duration: Option, +} + +/// A diagnostic record with its metadata and decoded payload. +#[derive(Clone, Debug, PartialEq, Eq, Serialize)] +#[serde(tag = "kind", rename_all = "lowercase")] +pub enum Diagnostic { + /// An operation began. + Start(DiagnosticStart), + /// An operation ended with a reported outcome and duration. + Finish(DiagnosticFinish), + /// A standalone observation. + Event(DiagnosticEvent), +} + +impl Diagnostic { + /// Returns the metadata shared by all diagnostic kinds. + pub fn key(&self) -> &DiagnosticKey { + match self { + Self::Start(start) => &start.key, + Self::Finish(finish) => &finish.key, + Self::Event(event) => &event.key, + } + } + + /// Returns whether this is a start, finish or standalone event. + pub fn kind(&self) -> Kind { + match self { + Self::Start(_) => Kind::Start, + Self::Finish(_) => Kind::Finish, + Self::Event(_) => Kind::Event, + } + } + + /// Returns the decoded payload; encoded data has already been decompressed. + pub fn payload(&self) -> &DiagnosticPayload { + match self { + Self::Start(start) => &start.payload, + Self::Finish(finish) => &finish.payload, + Self::Event(event) => &event.payload, + } + } +} + +/// A stored record that was not decoded as a diagnostic or report. +#[derive(Clone, Debug, PartialEq, Eq, Serialize)] +pub struct RawKeyValue { + /// Original key. + pub key: String, + /// Original value, without payload decoding. + pub value: String, + /// Why a recognized record could not be decoded; `None` for unknown or + /// non-diagnostic key/value pairs. + #[serde(skip_serializing_if = "Option::is_none")] + pub error: Option, +} + +/// One item returned by [`DiagnosticReader::entries`](crate::DiagnosticReader::entries). +#[derive(Clone, Debug, PartialEq, Eq, Serialize)] +#[serde(tag = "type")] +pub enum Entry { + /// A decoded diagnostic. + #[serde(rename = "diagnostic")] + Diagnostic(Diagnostic), + /// A provisioning health report. + #[serde(rename = "PROVISIONING_REPORT")] + Report(ProvisioningReport), + /// An unrelated or invalid record, preserved without interpretation. + #[serde(rename = "raw")] + Raw(RawKeyValue), +} + +fn serialize_timestamp( + timestamp: &DateTime, + serializer: S, +) -> Result +where + S: Serializer, +{ + serializer + .serialize_str(×tamp.to_rfc3339_opts(SecondsFormat::AutoSi, true)) +} + +fn serialize_encoding( + encoding: &Option, + serializer: S, +) -> Result +where + S: Serializer, +{ + match encoding { + Some(encoding) => encoding.serialize(serializer), + None => serializer.serialize_str("none"), + } +} + +fn serialize_duration_seconds( + duration: &Duration, + serializer: S, +) -> Result +where + S: Serializer, +{ + serializer.serialize_f64(duration.as_secs_f64()) +} + +fn serialize_opt_duration_seconds( + duration: &Option, + serializer: S, +) -> Result +where + S: Serializer, +{ + duration + .as_ref() + .map(Duration::as_secs_f64) + .serialize(serializer) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::ReportPpsType; + use rstest::rstest; + use serde_json::{json, Value}; + + const AGENT: &str = "azure-init/0.1.1"; + const VM_ID: &str = "3f2504e0-4f89-41d3-9a0c-0305e82c3301"; + const EVENT_ID: &str = "8f3e9c4a-1b2c-4d5e-9f01-234567890abc"; + const TIMESTAMP: &str = "2026-08-31T12:34:56.789Z"; + + fn key() -> DiagnosticKey { + DiagnosticKey { + agent: AGENT.into(), + vm_id: Some(VM_ID.into()), + name: "provision:run".into(), + event_id: EVENT_ID.into(), + timestamp: DateTime::parse_from_rfc3339(TIMESTAMP) + .unwrap() + .with_timezone(&Utc), + encoding: None, + } + } + + fn expected_diagnostic(kind: &str, payload: Value) -> Value { + json!({ + "type": "diagnostic", + "kind": kind, + "agent": AGENT, + "vm_id": VM_ID, + "name": "provision:run", + "event_id": EVENT_ID, + "timestamp": TIMESTAMP, + "encoding": "none", + "payload": payload, + }) + } + + #[rstest] + #[case(Kind::Start, "start")] + #[case(Kind::Finish, "finish")] + #[case(Kind::Event, "event")] + fn kind_uses_wire_token(#[case] kind: Kind, #[case] token: &str) { + assert_eq!(kind.to_string(), token); + assert_eq!(serde_json::to_value(kind).unwrap(), token); + } + + #[rstest] + #[case(Encoding::ZlibB64, "zlib+b64")] + #[case(Encoding::GzB64, "gz+b64")] + #[case(Encoding::Other("zstd+b64".into()), "zstd+b64")] + #[case(Encoding::Other("".into()), "")] + fn encoding_preserves_token( + #[case] encoding: Encoding, + #[case] token: &str, + ) { + assert_eq!(encoding.to_string(), token); + assert_eq!(serde_json::to_value(encoding).unwrap(), token); + } + + #[rstest] + #[case(Outcome::Success, "success")] + #[case(Outcome::Failure, "fail")] + fn outcome_uses_wire_token(#[case] result: Outcome, #[case] token: &str) { + assert_eq!(result.to_string(), token); + assert_eq!(serde_json::to_value(result).unwrap(), token); + } + + #[rstest] + #[case(DecodeError::UnsupportedVersion, "unsupported_version")] + #[case(DecodeError::IncompleteGroup, "incomplete_group")] + #[case(DecodeError::DuplicateChunk, "duplicate_chunk")] + #[case(DecodeError::Undecodable, "undecodable")] + #[case(DecodeError::Malformed, "malformed")] + fn decode_error_has_serializable_reason( + #[case] reason: DecodeError, + #[case] token: &str, + ) { + assert_eq!(serde_json::to_value(reason).unwrap(), token); + let error: &dyn std::error::Error = &reason; + assert!(!error.to_string().is_empty()); + assert!(error.source().is_none()); + } + + #[test] + fn string_conversions_preserve_text() { + let text = "héllo\n\"world\"\u{0}"; + let expected = DiagnosticPayload::Text(text.to_owned()); + assert_eq!(DiagnosticPayload::from(text), expected); + assert_eq!(DiagnosticPayload::from(text.to_owned()), expected); + } + + #[rstest] + #[case(b"")] + #[case(b"hello")] + #[case(&[0, 255, 128, 0])] + fn byte_inputs_remain_bytes(#[case] bytes: &[u8]) { + let expected = DiagnosticPayload::Bytes(bytes.to_vec()); + assert_eq!(DiagnosticPayload::from(bytes), expected); + assert_eq!(DiagnosticPayload::from(bytes.to_vec()), expected); + } + + #[rstest] + #[case("")] + #[case("héllo\n\"world\"\u{0}")] + fn text_serializes_as_a_string(#[case] text: &str) { + let payload = DiagnosticPayload::Text(text.to_owned()); + assert_eq!(serde_json::to_value(payload).unwrap(), json!(text)); + } + + #[rstest] + #[case(b"", "")] + #[case(b"hello", "aGVsbG8=")] + #[case(&[0], "AA==")] + #[case(&[251, 255], "+/8=")] + fn bytes_serialize_as_standard_base64( + #[case] bytes: &[u8], + #[case] encoded: &str, + ) { + assert_eq!( + serde_json::to_value(DiagnosticPayload::Bytes(bytes.to_vec())) + .unwrap(), + json!({"type": "bytes", "encoding": "base64", "data": encoded}) + ); + } + + #[test] + fn start_serializes_without_result_or_duration() { + let entry = Entry::Diagnostic(Diagnostic::Start(DiagnosticStart { + key: key(), + payload: "starting".into(), + })); + assert_eq!( + serde_json::to_value(entry).unwrap(), + expected_diagnostic("start", json!("starting")) + ); + } + + #[rstest] + #[case(Outcome::Success, "success", Duration::from_micros(312), 0.000312)] + #[case(Outcome::Failure, "fail", Duration::ZERO, 0.0)] + #[case(Outcome::Success, "success", Duration::MAX, u64::MAX as f64)] + fn finish_serializes_result_and_seconds( + #[case] result: Outcome, + #[case] token: &str, + #[case] duration: Duration, + #[case] seconds: f64, + ) { + let entry = Entry::Diagnostic(Diagnostic::Finish(DiagnosticFinish { + key: key(), + payload: "finished".into(), + result, + duration, + })); + let mut expected = expected_diagnostic("finish", json!("finished")); + expected["result"] = json!(token); + expected["duration"] = json!(seconds); + assert_eq!(serde_json::to_value(entry).unwrap(), expected); + } + + #[rstest] + #[case::neither(None, None)] + #[case::result_only(Some(Outcome::Success), None)] + #[case::zero_duration(None, Some(0))] + #[case::both(Some(Outcome::Failure), Some(52))] + fn event_serializes_only_measured_fields( + #[case] result: Option, + #[case] duration_us: Option, + ) { + let entry = Entry::Diagnostic(Diagnostic::Event(DiagnosticEvent { + key: key(), + payload: "observed".into(), + result, + duration: duration_us.map(Duration::from_micros), + })); + let mut expected = expected_diagnostic("event", json!("observed")); + if let Some(result) = result { + expected["result"] = json!(result.to_string()); + } + if let Some(duration_us) = duration_us { + expected["duration"] = json!(duration_us as f64 / 1_000_000.0); + } + assert_eq!(serde_json::to_value(entry).unwrap(), expected); + } + + #[test] + fn wire_encoding_is_distinct_from_payload_presentation() { + let entry = Entry::Diagnostic(Diagnostic::Event(DiagnosticEvent { + key: DiagnosticKey { + encoding: Some(Encoding::GzB64), + ..key() + }, + payload: b"hello".as_slice().into(), + result: None, + duration: None, + })); + let mut expected = expected_diagnostic( + "event", + json!({ + "type": "bytes", + "encoding": "base64", + "data": "aGVsbG8=", + }), + ); + expected["encoding"] = json!("gz+b64"); + assert_eq!(serde_json::to_value(entry).unwrap(), expected); + } + + #[rstest] + #[case(None, "none")] + #[case(Some(Encoding::ZlibB64), "zlib+b64")] + #[case(Some(Encoding::GzB64), "gz+b64")] + #[case(Some(Encoding::Other("zstd+b64".into())), "zstd+b64")] + fn key_always_serializes_encoding( + #[case] encoding: Option, + #[case] token: &str, + ) { + let key = DiagnosticKey { encoding, ..key() }; + assert_eq!(serde_json::to_value(key).unwrap()["encoding"], token); + } + + #[test] + fn missing_cloud_init_vm_id_serializes_as_null() { + let entry = Entry::Diagnostic(Diagnostic::Start(DiagnosticStart { + key: DiagnosticKey { + agent: "CLOUD_INIT".into(), + vm_id: None, + ..key() + }, + payload: "starting".into(), + })); + let mut expected = expected_diagnostic("start", json!("starting")); + expected["agent"] = json!("CLOUD_INIT"); + expected["vm_id"] = Value::Null; + assert_eq!(serde_json::to_value(entry).unwrap(), expected); + } + + #[rstest] + #[case("2026-08-31T12:34:56Z", "2026-08-31T12:34:56Z")] + #[case("2026-08-31T12:34:56.3Z", "2026-08-31T12:34:56.300Z")] + #[case("2026-08-31T12:34:56.789999Z", "2026-08-31T12:34:56.789999Z")] + #[case("2026-08-31T12:34:56.789123456Z", "2026-08-31T12:34:56.789123456Z")] + #[case("2026-08-31T14:34:56.789+02:00", TIMESTAMP)] + fn timestamp_serializes_in_utc_without_padding( + #[case] timestamp: &str, + #[case] expected: &str, + ) { + let key = DiagnosticKey { + timestamp: DateTime::parse_from_rfc3339(timestamp) + .unwrap() + .with_timezone(&Utc), + ..key() + }; + assert_eq!(serde_json::to_value(key).unwrap()["timestamp"], expected); + } + + #[rstest] + fn diagnostic_accessors_cover_all_kinds( + #[values(Kind::Start, Kind::Finish, Kind::Event)] kind: Kind, + ) { + let key = key(); + let payload = DiagnosticPayload::from("message"); + let diagnostic = match kind { + Kind::Start => Diagnostic::Start(DiagnosticStart { + key: key.clone(), + payload: payload.clone(), + }), + Kind::Finish => Diagnostic::Finish(DiagnosticFinish { + key: key.clone(), + payload: payload.clone(), + result: Outcome::Success, + duration: Duration::ZERO, + }), + Kind::Event => Diagnostic::Event(DiagnosticEvent { + key: key.clone(), + payload: payload.clone(), + result: None, + duration: None, + }), + }; + assert_eq!(diagnostic.key(), &key); + assert_eq!(diagnostic.kind(), kind); + assert_eq!(diagnostic.payload(), &payload); + } + + #[rstest] + #[case(None, None)] + #[case(Some(DecodeError::Malformed), Some("malformed"))] + fn raw_entry_preserves_key_value_and_optional_error( + #[case] error: Option, + #[case] token: Option<&str>, + ) { + let entry = Entry::Raw(RawKeyValue { + key: "original|key|0".into(), + value: "original \"value\"\nwith Unicode: é".into(), + error, + }); + let mut expected = json!({ + "type": "raw", + "key": "original|key|0", + "value": "original \"value\"\nwith Unicode: é", + }); + if let Some(token) = token { + expected["error"] = json!(token); + } + assert_eq!(serde_json::to_value(entry).unwrap(), expected); + } + + #[test] + fn report_entry_adds_type_without_nesting_report_fields() { + let report = + ProvisioningReport::success(AGENT, VM_ID, ReportPpsType::None); + let mut expected = serde_json::to_value(&report).unwrap(); + expected["type"] = json!("PROVISIONING_REPORT"); + assert_eq!( + serde_json::to_value(Entry::Report(report)).unwrap(), + expected + ); + } +} diff --git a/libazureinit-kvp/src/diagnostics/encoding.rs b/libazureinit-kvp/src/diagnostics/encoding.rs new file mode 100644 index 00000000..868cd4d2 --- /dev/null +++ b/libazureinit-kvp/src/diagnostics/encoding.rs @@ -0,0 +1,433 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +use std::io::{Read, Write}; + +use base64::{engine::general_purpose::STANDARD, Engine as _}; +use flate2::{ + bufread::{GzDecoder, ZlibDecoder}, + write::{GzEncoder, ZlibEncoder}, + Compression, +}; + +use super::diagnostic::{DecodeError, DiagnosticPayload, Encoding}; +use crate::KvpError; + +/// Produces a complete wire value; the writer handles chunk framing. +pub(super) fn encode_payload( + payload: DiagnosticPayload, + encoding: Option<&Encoding>, +) -> Result { + let bytes = match &payload { + DiagnosticPayload::Text(text) => text.as_bytes(), + DiagnosticPayload::Bytes(bytes) => bytes.as_slice(), + }; + match encoding { + None => { + let text = match payload { + DiagnosticPayload::Text(text) => text, + DiagnosticPayload::Bytes(bytes) => String::from_utf8(bytes) + .map_err(|_| KvpError::PayloadNotUtf8)?, + }; + if text.contains('\0') { + return Err(KvpError::ValueContainsNull); + } + Ok(text) + } + Some(Encoding::GzB64) => { + let mut encoder = + GzEncoder::new(Vec::new(), Compression::default()); + encoder.write_all(bytes)?; + Ok(STANDARD.encode(encoder.finish()?)) + } + Some(Encoding::ZlibB64) => { + let mut encoder = + ZlibEncoder::new(Vec::new(), Compression::default()); + encoder.write_all(bytes)?; + Ok(STANDARD.encode(encoder.finish()?)) + } + Some(Encoding::Other(token)) => Err(KvpError::UnsupportedEncoding { + token: token.clone(), + }), + } +} + +/// The reader must reassemble and validate chunk indices before decoding. +pub(super) fn decode_payload( + value: &[u8], + encoding: Option<&Encoding>, +) -> Result { + match encoding { + None => std::str::from_utf8(value) + .map(|text| DiagnosticPayload::Text(text.to_owned())) + .map_err(|_| DecodeError::Undecodable), + Some(encoding @ (Encoding::GzB64 | Encoding::ZlibB64)) => { + let compressed = STANDARD + .decode(value) + .map_err(|_| DecodeError::Undecodable)?; + decompress(&compressed, encoding) + } + Some(Encoding::Other(_)) => Err(DecodeError::Undecodable), + } +} + +pub(super) fn decompress( + compressed: &[u8], + encoding: &Encoding, +) -> Result { + let mut bytes = Vec::new(); + let remaining = match encoding { + Encoding::GzB64 => { + let mut decoder = GzDecoder::new(compressed); + decoder + .read_to_end(&mut bytes) + .map_err(|_| DecodeError::Undecodable)?; + decoder.into_inner() + } + Encoding::ZlibB64 => { + let mut decoder = ZlibDecoder::new(compressed); + decoder + .read_to_end(&mut bytes) + .map_err(|_| DecodeError::Undecodable)?; + decoder.into_inner() + } + Encoding::Other(_) => return Err(DecodeError::Undecodable), + }; + if !remaining.is_empty() { + return Err(DecodeError::Undecodable); + } + Ok(DiagnosticPayload::Bytes(bytes)) +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + // Python gzip fixtures keep decoding tests independent of our encoder. + const PYTHON_EMPTY: &str = "H4sIAAAAAAAC/wMAAAAAAAAAAAA="; + const PYTHON_HELLO: &str = "H4sIAAAAAAAC/8tIzcnJBwCGphA2BQAAAA=="; + const PYTHON_BINARY: &str = + "H4sIAAAAAAAC/0vOyS9N0c3MyyxRSMusKCktSuVi+A8AokCfWhUAAAA="; + const PYTHON_FILENAME: &str = + "H4sICAAAAAAC/2RtZXNnAEvOyS9N0c3MyyxRSMusKCktSuVi+A8AokCfWhUAAAA="; + const KUSTO_ZLIB: &str = "eJwLSS0uUSguKcrMS1cwNDIGACxqBQ4="; + + #[rstest] + #[case::empty_text("", false)] + #[case::empty_bytes("", true)] + #[case::unicode_text("héllo\n\"value\" | = 😀", false)] + #[case::unicode_bytes("héllo", true)] + #[case::not_inferred(PYTHON_HELLO, false)] + fn none_preserves_text(#[case] text: &str, #[case] as_bytes: bool) { + let payload = if as_bytes { + DiagnosticPayload::Bytes(text.as_bytes().to_vec()) + } else { + DiagnosticPayload::Text(text.to_owned()) + }; + let value = encode_payload(payload, None).unwrap(); + assert_eq!(value, text); + assert_eq!( + decode_payload(value.as_bytes(), None).unwrap(), + DiagnosticPayload::Text(text.to_owned()) + ); + } + + #[test] + fn none_keeps_large_payload_unencoded() { + let text = "é".repeat(4096); + let value = encode_payload(text.clone().into(), None).unwrap(); + assert_eq!(value, text); + assert_eq!( + decode_payload(value.as_bytes(), None).unwrap(), + DiagnosticPayload::Text(text) + ); + } + + #[rstest] + #[case(&[0xff])] + #[case(&[0xe2, 0x82])] + #[case(&[0xc0, 0xaf])] + #[case(&[0xed, 0xa0, 0x80])] + fn none_rejects_invalid_utf8(#[case] bytes: &[u8]) { + assert!(matches!( + encode_payload(bytes.into(), None), + Err(KvpError::PayloadNotUtf8) + )); + assert_eq!(decode_payload(bytes, None), Err(DecodeError::Undecodable)); + } + + #[rstest] + #[case::nul_only("\0", false)] + #[case::embedded_bytes("prefix\0suffix", true)] + #[case::trailing_text("trailing\0", false)] + fn none_rejects_nul_on_write(#[case] text: &str, #[case] as_bytes: bool) { + let payload = if as_bytes { + DiagnosticPayload::Bytes(text.as_bytes().to_vec()) + } else { + DiagnosticPayload::Text(text.to_owned()) + }; + assert!(matches!( + encode_payload(payload, None), + Err(KvpError::ValueContainsNull) + )); + } + + #[rstest] + #[case("")] + #[case("hello")] + #[case("héllo\n\0 | 😀")] + fn compressed_text_decodes_to_bytes( + #[case] text: &str, + #[values(Encoding::GzB64, Encoding::ZlibB64)] encoding: Encoding, + ) { + let value = encode_payload(text.into(), Some(&encoding)).unwrap(); + assert!(value.is_ascii()); + assert!(!value.contains('\0')); + assert_eq!( + decode_payload(value.as_bytes(), Some(&encoding)).unwrap(), + DiagnosticPayload::Bytes(text.as_bytes().to_vec()) + ); + } + + #[rstest] + #[case(vec![])] + #[case(b"hello".to_vec())] + #[case(vec![0xff, 0x80, 0, 0])] + #[case((0u8..=255).collect())] + fn compression_preserves_arbitrary_bytes( + #[case] bytes: Vec, + #[values(Encoding::GzB64, Encoding::ZlibB64)] encoding: Encoding, + ) { + let value = + encode_payload(bytes.clone().into(), Some(&encoding)).unwrap(); + assert_eq!( + decode_payload(value.as_bytes(), Some(&encoding)).unwrap(), + DiagnosticPayload::Bytes(bytes) + ); + } + + #[test] + fn gz_b64_writer_emits_gzip_header_crc_and_size() { + let value = + encode_payload("hello".into(), Some(&Encoding::GzB64)).unwrap(); + let gzip = STANDARD.decode(value).unwrap(); + assert!(gzip.starts_with(&[0x1f, 0x8b, 8, 0])); + assert_eq!( + &gzip[gzip.len() - 8..], + &[0x86, 0xa6, 0x10, 0x36, 5, 0, 0, 0] + ); + } + + #[test] + fn zlib_b64_matches_kusto_format() { + assert_eq!( + decode_payload(KUSTO_ZLIB.as_bytes(), Some(&Encoding::ZlibB64)) + .unwrap(), + DiagnosticPayload::Bytes(b"Test string 123".to_vec()), + ); + let value = + encode_payload("Test string 123".into(), Some(&Encoding::ZlibB64)) + .unwrap(); + let compressed = STANDARD.decode(value).unwrap(); + assert_eq!(compressed[0], 0x78); + assert_eq!(compressed[1] & 0x20, 0); + } + + #[test] + fn zlib_b64_requires_a_complete_valid_zlib_stream() { + let compressed = STANDARD.decode(KUSTO_ZLIB).unwrap(); + let mut corrupt = compressed.clone(); + *corrupt.last_mut().unwrap() ^= 1; + let mut trailing = compressed.clone(); + trailing.push(0); + for bytes in [ + STANDARD.decode(PYTHON_HELLO).unwrap(), + compressed[..compressed.len() - 1].to_vec(), + corrupt, + trailing, + ] { + assert_eq!( + decode_payload( + STANDARD.encode(bytes).as_bytes(), + Some(&Encoding::ZlibB64) + ), + Err(DecodeError::Undecodable), + ); + } + } + + #[rstest] + #[case::empty(PYTHON_EMPTY, b"")] + #[case::text(PYTHON_HELLO, b"hello")] + #[case::binary(PYTHON_BINARY, b"cloud-init fixture\n\x00\xff")] + #[case::filename(PYTHON_FILENAME, b"cloud-init fixture\n\x00\xff")] + fn decodes_python_gzip_fixtures( + #[case] value: &str, + #[case] expected: &[u8], + ) { + assert_eq!( + decode_payload(value.as_bytes(), Some(&Encoding::GzB64)).unwrap(), + DiagnosticPayload::Bytes(expected.to_vec()) + ); + } + + #[rstest] + #[case("!")] + #[case("====")] + #[case("AA=A")] + #[case("é")] + fn gz_b64_rejects_invalid_base64(#[case] value: &str) { + assert_eq!( + decode_payload(value.as_bytes(), Some(&Encoding::GzB64)), + Err(DecodeError::Undecodable) + ); + } + + #[test] + fn gz_b64_requires_canonical_base64() { + for value in [ + format!("{PYTHON_HELLO}\n"), + PYTHON_HELLO.trim_end_matches('=').to_owned(), + PYTHON_HELLO.replace("AA==", "AB=="), + ] { + assert_eq!( + decode_payload(value.as_bytes(), Some(&Encoding::GzB64)), + Err(DecodeError::Undecodable), + "{value}" + ); + } + } + + #[rstest] + #[case(b"")] + #[case(b"not gzip")] + #[case(b"\x78\x9c\x03\x00\x00\x00\x00\x01")] + fn gz_b64_rejects_non_gzip_data(#[case] bytes: &[u8]) { + let value = STANDARD.encode(bytes); + assert_eq!( + decode_payload(value.as_bytes(), Some(&Encoding::GzB64)), + Err(DecodeError::Undecodable) + ); + } + + #[rstest] + #[case(PYTHON_EMPTY)] + #[case(PYTHON_HELLO)] + #[case(PYTHON_BINARY)] + #[case(PYTHON_FILENAME)] + fn every_truncated_gzip_prefix_is_undecodable(#[case] value: &str) { + let gzip = STANDARD.decode(value).unwrap(); + for end in 0..gzip.len() { + let truncated = STANDARD.encode(&gzip[..end]); + assert_eq!( + decode_payload(truncated.as_bytes(), Some(&Encoding::GzB64)), + Err(DecodeError::Undecodable), + "gzip truncated at byte {end}" + ); + } + } + + #[test] + fn every_truncated_base64_prefix_is_undecodable() { + for end in 0..PYTHON_HELLO.len() { + assert_eq!( + decode_payload( + &PYTHON_HELLO.as_bytes()[..end], + Some(&Encoding::GzB64) + ), + Err(DecodeError::Undecodable), + "base64 truncated at byte {end}" + ); + } + } + + #[rstest] + #[case::magic(0)] + #[case::compression_method(2)] + #[case::deflate_body(10)] + fn gz_b64_rejects_corrupt_gzip(#[case] offset: usize) { + let mut gzip = STANDARD.decode(PYTHON_HELLO).unwrap(); + gzip[offset] ^= 1; + let value = STANDARD.encode(gzip); + assert_eq!( + decode_payload(value.as_bytes(), Some(&Encoding::GzB64)), + Err(DecodeError::Undecodable) + ); + } + + #[rstest] + #[case::crc(8)] + #[case::size(4)] + fn gz_b64_validates_crc_and_size(#[case] trailer_offset: usize) { + let mut gzip = STANDARD.decode(PYTHON_HELLO).unwrap(); + let offset = gzip.len() - trailer_offset; + gzip[offset] ^= 1; + let value = STANDARD.encode(gzip); + assert_eq!( + decode_payload(value.as_bytes(), Some(&Encoding::GzB64)), + Err(DecodeError::Undecodable) + ); + } + + #[rstest] + #[case(b"\0")] + #[case(b"trailing data")] + fn gz_b64_rejects_trailing_data(#[case] suffix: &[u8]) { + let mut gzip = STANDARD.decode(PYTHON_HELLO).unwrap(); + gzip.extend_from_slice(suffix); + let value = STANDARD.encode(gzip); + assert_eq!( + decode_payload(value.as_bytes(), Some(&Encoding::GzB64)), + Err(DecodeError::Undecodable) + ); + } + + #[rstest] + #[case(PYTHON_EMPTY)] + #[case(PYTHON_HELLO)] + fn gz_b64_rejects_concatenated_members(#[case] second: &str) { + let mut gzip = STANDARD.decode(PYTHON_HELLO).unwrap(); + gzip.extend(STANDARD.decode(second).unwrap()); + let value = STANDARD.encode(gzip); + assert_eq!( + decode_payload(value.as_bytes(), Some(&Encoding::GzB64)), + Err(DecodeError::Undecodable) + ); + } + + #[test] + fn gz_b64_handles_large_expansion() { + let bytes = vec![0; 2 * 1024 * 1024]; + let encoding = Some(&Encoding::GzB64); + let value = encode_payload(bytes.clone().into(), encoding).unwrap(); + assert!(value.len() < bytes.len()); + assert_eq!( + decode_payload(value.as_bytes(), encoding).unwrap(), + DiagnosticPayload::Bytes(bytes) + ); + } + + #[rstest] + #[case::unknown("zstd+b64")] + #[case::plain_token("none")] + #[case::gzip_token("gz+b64")] + #[case::zlib_token("zlib+b64")] + fn other_encoding_is_never_inferred(#[case] token: &str) { + let encoding = Encoding::Other(token.into()); + assert!(matches!( + encode_payload("payload".into(), Some(&encoding)), + Err(KvpError::UnsupportedEncoding { token: actual }) + if actual == token + )); + assert_eq!( + decode_payload(PYTHON_HELLO.as_bytes(), Some(&encoding)), + Err(DecodeError::Undecodable) + ); + let compressed = STANDARD.decode(PYTHON_HELLO).unwrap(); + assert_eq!( + decompress(&compressed, &encoding), + Err(DecodeError::Undecodable) + ); + } +} diff --git a/libazureinit-kvp/src/diagnostics/mod.rs b/libazureinit-kvp/src/diagnostics/mod.rs new file mode 100644 index 00000000..293441d7 --- /dev/null +++ b/libazureinit-kvp/src/diagnostics/mod.rs @@ -0,0 +1,31 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +//! Diagnostic types and KVP pool access. + +mod cloud_init; +mod diagnostic; +mod encoding; +mod reader; +mod writer; + +pub use diagnostic::{ + DecodeError, Diagnostic, DiagnosticEvent, DiagnosticFinish, DiagnosticKey, + DiagnosticPayload, DiagnosticStart, Encoding, Entry, Kind, Outcome, + RawKeyValue, DIAGNOSTIC_VERSION_ID, +}; +pub use reader::DiagnosticReader; +pub use writer::{DiagnosticWriter, DurationPrecision, TimestampPrecision}; + +/// Maximum encoded payload bytes stored in one diagnostic record. +/// +/// [`DiagnosticWriter`] splits larger payloads automatically. +pub const MAX_CHUNK_BYTES: usize = 1022; + +/// Parse a non-empty run of ASCII digits (a chunk index or duration) as `u64`. +fn parse_unsigned(value: &str) -> Result { + if value.is_empty() || !value.bytes().all(|byte| byte.is_ascii_digit()) { + return Err(DecodeError::Malformed); + } + value.parse().map_err(|_| DecodeError::Malformed) +} diff --git a/libazureinit-kvp/src/diagnostics/reader.rs b/libazureinit-kvp/src/diagnostics/reader.rs new file mode 100644 index 00000000..c2cd7fd6 --- /dev/null +++ b/libazureinit-kvp/src/diagnostics/reader.rs @@ -0,0 +1,1158 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +use std::collections::HashMap; +use std::time::Duration; + +use chrono::{DateTime, Utc}; +use uuid::Uuid; + +use super::cloud_init; +use super::diagnostic::{ + DecodeError, Diagnostic, DiagnosticEvent, DiagnosticFinish, DiagnosticKey, + DiagnosticPayload, DiagnosticStart, Encoding, Entry, Outcome, RawKeyValue, + DIAGNOSTIC_VERSION_ID, +}; +use super::encoding::decode_payload; +use super::parse_unsigned; +use crate::{ + KvpError, KvpPoolStore, ProvisioningReport, PROVISIONING_REPORT_KEY, +}; + +/// Reads diagnostics and provisioning reports from a KVP pool. +/// +/// Native and cloud-init diagnostics are decoded. Other or invalid records are +/// returned as [`Entry::Raw`]. Agent and VM identities come from the records, +/// so no local identity is required. +#[derive(Clone, Debug)] +pub struct DiagnosticReader { + store: KvpPoolStore, +} + +impl DiagnosticReader { + /// Creates a reader for `store`; the pool is accessed by [`entries`](Self::entries). + pub fn new(store: KvpPoolStore) -> Self { + Self { store } + } + + /// Reads the current pool contents without modifying them. + /// + /// Returns an empty list if the pool file does not exist. Entries retain + /// pool order, and unrelated or invalid records remain [`Entry::Raw`]. + /// + /// # Errors + /// Returns [`KvpError`] if the pool cannot be read or has invalid storage + /// layout or UTF-8; no entries are returned. Per-record decoding failures + /// are returned as [`Entry::Raw`] values with a [`DecodeError`]. + pub fn entries(&self) -> Result, KvpError> { + Ok(decode_entries(self.store.dump()?)) + } +} + +struct Chunk { + position: usize, + index: u64, + raw: RawKeyValue, +} + +struct ChunkGroup { + first_position: usize, + chunks: Vec, +} + +fn decode_entries(records: Vec<(String, String)>) -> Vec { + let mut entries = Vec::new(); + let mut groups = HashMap::::new(); + + for (position, (key, value)) in records.into_iter().enumerate() { + let mut raw = RawKeyValue { + key, + value, + error: None, + }; + let prefix = raw.key.split('|').next().unwrap_or_default(); + let parsed = match prefix { + DIAGNOSTIC_VERSION_ID => split_chunk_key(&raw.key) + .map(|(base, index)| (base, Some(index))), + cloud_init::PREFIX => cloud_init::split_key(&raw.key), + PROVISIONING_REPORT_KEY if raw.key == PROVISIONING_REPORT_KEY => { + match raw.value.parse::() { + Ok(report) => { + entries.push((position, Entry::Report(report))) + } + Err(error) => { + raw.error = Some(error); + entries.push((position, Entry::Raw(raw))); + } + } + continue; + } + version if version.starts_with("DIAG_V") => { + Err(DecodeError::UnsupportedVersion) + } + _ => { + entries.push((position, Entry::Raw(raw))); + continue; + } + }; + match parsed { + Ok((base, Some(index))) => { + groups + .entry(base.to_owned()) + .or_insert_with(|| ChunkGroup { + first_position: position, + chunks: Vec::new(), + }) + .chunks + .push(Chunk { + position, + index, + raw, + }); + } + Ok((base, None)) => { + match cloud_init::decode_single(base, &raw.value) { + Ok(diagnostic) => { + entries.push((position, Entry::Diagnostic(diagnostic))) + } + Err(error) => { + raw.error = Some(error); + entries.push((position, Entry::Raw(raw))); + } + } + } + Err(error) => { + raw.error = Some(error); + entries.push((position, Entry::Raw(raw))); + } + } + } + + for (base, mut group) in groups { + let decoded = if base.split('|').next() == Some(cloud_init::PREFIX) { + order_chunks(&mut group.chunks).and_then(|()| { + cloud_init::decode_chunks( + &base, + group + .chunks + .iter() + .map(|chunk| (chunk.index, chunk.raw.value.as_str())), + ) + }) + } else { + decode_diag_group(&base, &mut group.chunks) + }; + match decoded { + Ok(diagnostic) => { + entries.push(( + group.first_position, + Entry::Diagnostic(diagnostic), + )); + } + Err(error) => { + for mut chunk in group.chunks { + chunk.raw.error = Some(error); + entries.push((chunk.position, Entry::Raw(chunk.raw))); + } + } + } + } + + // Failed groups retain every physical record at its original position. + entries.sort_by_key(|(position, _)| *position); + entries.into_iter().map(|(_, entry)| entry).collect() +} + +fn split_chunk_key(key: &str) -> Result<(&str, u64), DecodeError> { + let (base, index) = key.rsplit_once('|').ok_or(DecodeError::Malformed)?; + Ok((base, parse_unsigned(index)?)) +} + +fn decode_diag_group( + base: &str, + chunks: &mut [Chunk], +) -> Result { + let [version, agent, vm_id, kind, name, event_id, timestamp, encoding, result, duration]: [&str; 10] = + base.split('|') + .collect::>() + .try_into() + .map_err(|_| DecodeError::Malformed)?; + if version != DIAGNOSTIC_VERSION_ID + || base.contains('\0') + || [agent, vm_id, name, event_id, timestamp, encoding] + .iter() + .any(|field| field.is_empty()) + { + return Err(DecodeError::Malformed); + } + Uuid::parse_str(vm_id).map_err(|_| DecodeError::Malformed)?; + let parsed_timestamp = DateTime::parse_from_rfc3339(timestamp) + .map_err(|_| DecodeError::Malformed)? + .with_timezone(&Utc); + let result = match result { + "" => None, + "success" => Some(Outcome::Success), + "fail" => Some(Outcome::Failure), + _ => return Err(DecodeError::Malformed), + }; + let duration = if duration.is_empty() { + None + } else { + Some(parse_duration(duration)?) + }; + let encoding = match encoding { + "none" => None, + "zlib+b64" => Some(Encoding::ZlibB64), + "gz+b64" => Some(Encoding::GzB64), + other => Some(Encoding::Other(other.to_owned())), + }; + let key = DiagnosticKey { + agent: agent.to_owned(), + vm_id: Some(vm_id.to_owned()), + name: name.to_owned(), + event_id: event_id.to_owned(), + timestamp: parsed_timestamp, + encoding, + }; + + match (kind, result, duration) { + ("start", None, None) => { + let payload = decode_chunks(chunks, key.encoding.as_ref())?; + Ok(Diagnostic::Start(DiagnosticStart { key, payload })) + } + ("finish", Some(result), Some(duration)) => { + let payload = decode_chunks(chunks, key.encoding.as_ref())?; + Ok(Diagnostic::Finish(DiagnosticFinish { + key, + payload, + result, + duration, + })) + } + ("event", result, duration) => { + let payload = decode_chunks(chunks, key.encoding.as_ref())?; + Ok(Diagnostic::Event(DiagnosticEvent { + key, + payload, + result, + duration, + })) + } + _ => Err(DecodeError::Malformed), + } +} + +fn parse_duration(value: &str) -> Result { + parse_decimal_duration(value).or_else(|_| { + let seconds = + value.parse::().map_err(|_| DecodeError::Malformed)?; + Duration::try_from_secs_f64(seconds).map_err(|_| DecodeError::Malformed) + }) +} + +fn parse_decimal_duration(value: &str) -> Result { + let (seconds, fraction) = value + .split_once('.') + .map_or((value, None), |(seconds, fraction)| { + (seconds, Some(fraction)) + }); + let seconds = parse_unsigned(seconds)?; + let nanos = match fraction { + None => 0, + Some(fraction) => { + if fraction.is_empty() + || fraction.len() > 9 + || !fraction.bytes().all(|byte| byte.is_ascii_digit()) + { + return Err(DecodeError::Malformed); + } + fraction + .parse::() + .map_err(|_| DecodeError::Malformed)? + * 10u32.pow(9 - fraction.len() as u32) + } + }; + Ok(Duration::new(seconds, nanos)) +} + +fn decode_chunks( + chunks: &mut [Chunk], + encoding: Option<&Encoding>, +) -> Result { + order_chunks(chunks)?; + let value: String = chunks + .iter() + .map(|chunk| chunk.raw.value.as_str()) + .collect(); + decode_payload(value.as_bytes(), encoding) +} + +fn order_chunks(chunks: &mut [Chunk]) -> Result<(), DecodeError> { + chunks.sort_by_key(|chunk| chunk.index); + if chunks.windows(2).any(|pair| pair[0].index == pair[1].index) { + return Err(DecodeError::DuplicateChunk); + } + if chunks + .iter() + .enumerate() + .any(|(expected, chunk)| usize::try_from(chunk.index) != Ok(expected)) + { + return Err(DecodeError::IncompleteGroup); + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + use std::fs; + use std::io; + use std::path::Path; + use std::sync::atomic::{AtomicUsize, Ordering}; + use std::sync::Arc; + + use base64::{engine::general_purpose::STANDARD, Engine as _}; + use chrono::SecondsFormat; + use rstest::rstest; + use tempfile::TempDir; + + use super::super::diagnostic::Kind; + use crate::store::{Handle, OsSysOps, StatInfo, SysOps}; + use crate::{write_report, KvpPool, PoolMode, ReportPpsType}; + + const AGENT: &str = "azure-init/0.1.1"; + const VM_ID: &str = "3f2504e0-4f89-41d3-9a0c-0305e82c3301"; + const EVENT_ID: &str = "8f3e9c4a-1b2c-4d5e-9f01-234567890abc"; + const TIMESTAMP: &str = "2026-08-31T12:34:56.789Z"; + const GZIP_HELLO: &str = "H4sIAAAAAAAC/8tIzcnJBwCGphA2BQAAAA=="; + + fn key(index: u64) -> String { + format!( + "DIAG|{AGENT}|{VM_ID}|event|test|{EVENT_ID}|{TIMESTAMP}|none|||{index}" + ) + } + + fn with_field(key: &str, index: usize, value: &str) -> String { + let mut fields: Vec<_> = key.split('|').collect(); + fields[index] = value; + fields.join("|") + } + + fn raw_entries( + records: &[(String, String)], + error: Option, + ) -> Vec { + records + .iter() + .map(|(key, value)| { + Entry::Raw(RawKeyValue { + key: key.clone(), + value: value.clone(), + error, + }) + }) + .collect() + } + + fn only_diagnostic(mut entries: Vec) -> Diagnostic { + assert_eq!(entries.len(), 1); + let mut diagnostic = None; + if let Some(Entry::Diagnostic(value)) = entries.pop() { + diagnostic = Some(value); + } + diagnostic.expect("expected a diagnostic") + } + + fn store(dir: &TempDir) -> KvpPoolStore { + KvpPoolStore::new_in(KvpPool::Guest, dir.path(), PoolMode::Safe) + .unwrap() + } + + #[derive(Debug, Default)] + struct ReaderOps { + os: OsSysOps, + calls: AtomicUsize, + open_error: Option, + } + + impl SysOps for ReaderOps { + fn open_read(&self, path: &Path) -> io::Result> { + self.calls.fetch_add(1, Ordering::SeqCst); + if let Some(error) = self.open_error { + return Err(error.into()); + } + self.os.open_read(path) + } + + fn open_read_write(&self, _: &Path) -> io::Result> { + self.calls.fetch_add(1, Ordering::SeqCst); + Err(io::ErrorKind::Unsupported.into()) + } + + fn open_read_write_create( + &self, + _: &Path, + ) -> io::Result> { + self.calls.fetch_add(1, Ordering::SeqCst); + Err(io::ErrorKind::Unsupported.into()) + } + + fn path_metadata(&self, _: &Path) -> io::Result { + self.calls.fetch_add(1, Ordering::SeqCst); + Err(io::ErrorKind::Unsupported.into()) + } + + fn boot_time(&self) -> io::Result { + self.calls.fetch_add(1, Ordering::SeqCst); + Err(io::ErrorKind::Unsupported.into()) + } + } + + fn observed_reader(dir: &TempDir) -> (DiagnosticReader, Arc) { + let ops = Arc::new(ReaderOps::default()); + let observed = KvpPoolStore::with_ops( + KvpPool::Guest, + dir.path(), + PoolMode::Safe, + ops.clone(), + ) + .unwrap(); + (DiagnosticReader::new(observed), ops) + } + + #[test] + fn reader_ops_rejects_non_read_operations() { + let dir = TempDir::new().unwrap(); + let pool = store(&dir); + let ops = ReaderOps::default(); + assert_eq!( + [ + ops.open_read_write(pool.path()).unwrap_err().kind(), + ops.open_read_write_create(pool.path()).unwrap_err().kind(), + ops.path_metadata(pool.path()).unwrap_err().kind(), + ops.boot_time().unwrap_err().kind(), + ], + [io::ErrorKind::Unsupported; 4] + ); + assert_eq!(ops.calls.load(Ordering::SeqCst), 4); + assert!(!pool.path().exists()); + } + + #[test] + fn constructor_does_no_io() { + let dir = TempDir::new().unwrap(); + let (reader, ops) = observed_reader(&dir); + assert_eq!(ops.calls.load(Ordering::SeqCst), 0); + assert!(!reader.store.path().exists()); + } + + #[test] + fn entries_reads_one_fresh_snapshot_without_modifying_the_pool() { + let dir = TempDir::new().unwrap(); + let pool = store(&dir); + let (reader, ops) = observed_reader(&dir); + assert_eq!(ops.calls.load(Ordering::SeqCst), 0); + assert!(reader.entries().unwrap().is_empty()); + assert_eq!(ops.calls.load(Ordering::SeqCst), 1); + assert!(!pool.path().exists()); + + pool.append("unrelated", "unchanged").unwrap(); + let before = fs::read(pool.path()).unwrap(); + assert_eq!( + reader.entries().unwrap(), + raw_entries(&[("unrelated".into(), "unchanged".into())], None) + ); + assert_eq!(ops.calls.load(Ordering::SeqCst), 2); + assert_eq!(fs::read(pool.path()).unwrap(), before); + } + + #[test] + fn snapshot_open_errors_propagate() { + let dir = TempDir::new().unwrap(); + let ops = Arc::new(ReaderOps { + open_error: Some(io::ErrorKind::PermissionDenied), + ..ReaderOps::default() + }); + let pool = KvpPoolStore::with_ops( + KvpPool::Guest, + dir.path(), + PoolMode::Safe, + ops.clone(), + ) + .unwrap(); + let reader = DiagnosticReader::new(pool); + assert_eq!(ops.calls.load(Ordering::SeqCst), 0); + assert!(matches!( + reader.entries(), + Err(KvpError::Io(error)) + if error.kind() == io::ErrorKind::PermissionDenied + )); + assert_eq!(ops.calls.load(Ordering::SeqCst), 1); + } + + #[test] + fn truncated_physical_pool_is_a_snapshot_error() { + let dir = TempDir::new().unwrap(); + let pool = store(&dir); + fs::write(pool.path(), b"partial record").unwrap(); + let reader = DiagnosticReader::new(pool); + assert!(matches!( + reader.entries(), + Err(KvpError::Io(error)) if error.kind() == io::ErrorKind::Other + )); + } + + #[rstest] + #[case::key(0)] + #[case::value(512)] + fn non_utf8_record_fails_the_entire_snapshot(#[case] offset: usize) { + let dir = TempDir::new().unwrap(); + let pool = store(&dir); + pool.append(&with_field(&key(0), 4, "before"), "valid") + .unwrap(); + let record_start = fs::read(pool.path()).unwrap().len(); + pool.append(&key(0), "value").unwrap(); + pool.append(&with_field(&key(0), 4, "after"), "valid") + .unwrap(); + let reader = DiagnosticReader::new(pool.clone()); + assert_eq!(reader.entries().unwrap().len(), 3); + + let mut bytes = fs::read(pool.path()).unwrap(); + bytes[record_start + offset] = 0xff; + fs::write(pool.path(), &bytes).unwrap(); + assert!(matches!( + reader.entries(), + Err(KvpError::Io(error)) if error.kind() == io::ErrorKind::InvalidData + )); + assert_eq!(fs::read(pool.path()).unwrap(), bytes); + } + + #[rstest] + #[case("")] + #[case("unrelated")] + #[case("other|0")] + #[case("DIAG_OTHER|0")] + #[case("diag|0")] + #[case("prefixDIAG_V2|0")] + #[case("azure-init-0.1.1|1700000000|vm-abc|event|imds|id|2026-08-31T12:34:56.789Z|0")] + fn unrelated_and_pre_adoption_records_remain_raw(#[case] key: &str) { + let records = vec![ + (key.into(), "first\nvalue".into()), + (key.into(), "second \"value\"".into()), + ]; + assert_eq!( + decode_entries(records.clone()), + raw_entries(&records, None) + ); + } + + #[rstest] + #[case::missing_fields("result=success".into())] + #[case::conflicting_results(format!( + "result=success|agent={AGENT}|pps_type=None|vm_id={VM_ID}|timestamp={TIMESTAMP}|result=error" + ))] + fn malformed_report_is_preserved_without_losing_other_entries( + #[case] value: String, + ) { + let records = vec![(PROVISIONING_REPORT_KEY.into(), value)]; + let mut mixed = records.clone(); + mixed.push((key(0), "valid diagnostic".into())); + let entries = decode_entries(mixed); + assert_eq!(entries.len(), 2); + assert_eq!( + entries[0], + raw_entries(&records, Some(DecodeError::Malformed))[0] + ); + assert!(matches!(entries[1], Entry::Diagnostic(_))); + } + + #[test] + fn report_is_typed_in_first_seen_order() { + let dir = TempDir::new().unwrap(); + let pool = store(&dir); + let report = ProvisioningReport::failure( + AGENT, + VM_ID, + "failed | with details", + ReportPpsType::None, + ) + .with_extra("detail", "first") + .with_extra("detail", "second"); + pool.append(&key(1), "second chunk").unwrap(); + write_report(&pool, &report).unwrap(); + pool.append(&key(0), "first chunk").unwrap(); + let before = fs::read(pool.path()).unwrap(); + let entries = DiagnosticReader::new(pool.clone()).entries().unwrap(); + assert_eq!(entries.len(), 2); + assert!(matches!(entries[0], Entry::Diagnostic(_))); + assert_eq!(entries[1], Entry::Report(report)); + assert_eq!(fs::read(pool.path()).unwrap(), before); + } + + #[test] + fn repeated_report_records_are_not_deduplicated() { + let dir = TempDir::new().unwrap(); + let pool = store(&dir); + let report = + ProvisioningReport::success(AGENT, VM_ID, ReportPpsType::None); + write_report(&pool, &report).unwrap(); + let value = pool.read(PROVISIONING_REPORT_KEY).unwrap().unwrap(); + pool.append(PROVISIONING_REPORT_KEY, &value).unwrap(); + assert_eq!( + DiagnosticReader::new(pool).entries().unwrap(), + vec![Entry::Report(report.clone()), Entry::Report(report)] + ); + } + + #[test] + fn report_prefix_is_not_a_chunked_report_key() { + let records = vec![("PROVISIONING_REPORT|0".into(), "value".into())]; + assert_eq!( + decode_entries(records.clone()), + raw_entries(&records, None) + ); + } + + #[rstest] + #[case("DIAG_V0")] + #[case("DIAG_V1")] + #[case("DIAG_V2")] + #[case("DIAG_V999")] + #[case("DIAG_V1_extra")] + #[case("DIAG_V")] + fn unsupported_versions_do_not_interpret_later_fields( + #[case] version: &str, + ) { + let records = vec![ + (version.into(), "value".into()), + (format!("{version}|bad metadata|0"), "first".into()), + (format!("{version}|bad metadata|0"), "duplicate".into()), + (format!("{version}|bad metadata|2"), "gap".into()), + ]; + assert_eq!( + decode_entries(records.clone()), + raw_entries(&records, Some(DecodeError::UnsupportedVersion)) + ); + } + + #[rstest] + #[case::uuid(EVENT_ID)] + #[case::ten_digits("0000000001")] + #[case::opaque("bad-uuid")] + fn native_event_decodes_without_local_identity(#[case] event_id: &str) { + let payload = "héllo\n\"message\" | ="; + let diagnostic = only_diagnostic(decode_entries(vec![( + with_field(&key(0), 5, event_id), + payload.into(), + )])); + assert_eq!(diagnostic.kind(), Kind::Event); + assert_eq!(diagnostic.key().agent, AGENT); + assert_eq!(diagnostic.key().vm_id.as_deref(), Some(VM_ID)); + assert_eq!(diagnostic.key().name, "test"); + assert_eq!(diagnostic.key().event_id, event_id); + assert_eq!(diagnostic.key().encoding, None); + assert_eq!( + diagnostic + .key() + .timestamp + .to_rfc3339_opts(SecondsFormat::Millis, true), + TIMESTAMP + ); + assert_eq!( + diagnostic.payload(), + &DiagnosticPayload::Text(payload.into()) + ); + } + + #[rstest] + #[case::unmatched_start("start", "", "")] + #[case::orphan_finish("finish", "fail", "0.000312")] + fn isolated_span_endpoint_decodes( + #[case] kind: &str, + #[case] result: &str, + #[case] duration: &str, + ) { + let key = with_field(&key(0), 3, kind); + let key = with_field(&key, 8, result); + let key = with_field(&key, 9, duration); + let diagnostic = + only_diagnostic(decode_entries(vec![(key, "message".into())])); + assert_eq!(diagnostic.kind().to_string(), kind); + } + + #[test] + fn span_endpoints_with_a_shared_id_decode_separately() { + let start = with_field(&key(0), 3, "start"); + let finish = with_field(&key(0), 3, "finish"); + let finish = with_field(&finish, 8, "fail"); + let finish = with_field(&finish, 9, "0.000312"); + let entries = decode_entries(vec![ + (start, "starting".into()), + (finish, "failed".into()), + ]); + assert!(matches!( + entries.as_slice(), + [Entry::Diagnostic(Diagnostic::Start(start)), Entry::Diagnostic(Diagnostic::Finish(finish))] + if start.key.event_id == EVENT_ID + && finish.key.event_id == EVENT_ID + && start.payload == DiagnosticPayload::Text("starting".into()) + && finish.payload == DiagnosticPayload::Text("failed".into()) + && finish.result == Outcome::Failure + && finish.duration == Duration::from_micros(312) + )); + } + + #[rstest] + #[case::neither(None, None)] + #[case::result_only(Some(Outcome::Success), None)] + #[case::zero_duration(None, Some(0))] + #[case::both(Some(Outcome::Failure), Some(52))] + fn event_result_and_duration_are_independent( + #[case] result: Option, + #[case] duration_secs: Option, + ) { + let result_token = result.map_or_else(String::new, |v| v.to_string()); + let duration_token = + duration_secs.map_or_else(String::new, |value| value.to_string()); + let key = with_field(&key(0), 8, &result_token); + let key = with_field(&key, 9, &duration_token); + let diagnostic = + only_diagnostic(decode_entries(vec![(key, "value".into())])); + assert!(matches!(&diagnostic, Diagnostic::Event(event) + if event.result == result + && event.duration == duration_secs.map(Duration::from_secs))); + } + + #[rstest] + #[case("start", "success", "")] + #[case("start", "", "0")] + #[case("start", "fail", "52")] + #[case("finish", "", "")] + #[case("finish", "success", "")] + #[case("finish", "", "0")] + #[case("finish", "error", "52")] + fn invalid_kind_fields_are_malformed( + #[case] kind: &str, + #[case] result: &str, + #[case] duration: &str, + ) { + let base = with_field(&key(0), 3, kind); + let base = with_field(&base, 8, result); + let base = with_field(&base, 9, duration); + let records = vec![(base, "value".into())]; + assert_eq!( + decode_entries(records.clone()), + raw_entries(&records, Some(DecodeError::Malformed)) + ); + } + + #[rstest] + #[case::empty_agent(1, "")] + #[case::invalid_vm_id(2, "vm-abc")] + #[case::invalid_kind(3, "compressed")] + #[case::empty_name(4, "")] + #[case::null_in_name(4, "bad\0name")] + #[case::empty_event_id(5, "")] + #[case::empty_encoding(7, "")] + #[case::invalid_result(8, "SUCCESS")] + #[case::negative_duration(9, "-0.1")] + #[case::non_finite_duration(9, "NaN")] + #[case::infinite_duration(9, "inf")] + #[case::duration_whitespace(9, " 1")] + #[case::float_overflow(9, "1e100")] + #[case::signed_fraction(9, "1.+2")] + #[case::numeric_overflow(9, "18446744073709551616")] + #[case::missing_chunk_index(10, "")] + #[case::non_numeric_chunk_index(10, "x")] + fn malformed_native_fields_preserve_the_original_record( + #[case] field: usize, + #[case] value: &str, + ) { + let records = vec![(with_field(&key(0), field, value), "value".into())]; + assert_eq!( + decode_entries(records.clone()), + raw_entries(&records, Some(DecodeError::Malformed)) + ); + } + + #[test] + fn native_format_requires_the_exact_key_layout_and_an_index() { + let complete = key(0); + let records = vec![ + ("DIAG".into(), "value".into()), + (complete.rsplit_once('|').unwrap().0.into(), "value".into()), + (format!("{complete}|1"), "value".into()), + ]; + assert_eq!( + decode_entries(records.clone()), + raw_entries(&records, Some(DecodeError::Malformed)) + ); + } + + #[rstest] + #[case::unparsable("not a timestamp")] + #[case::missing_offset("2026-08-31T12:34:56.789")] + #[case::invalid_month("2026-13-31T12:34:56.789Z")] + fn native_format_rejects_invalid_timestamps(#[case] timestamp: &str) { + let records = vec![(with_field(&key(0), 6, timestamp), "value".into())]; + assert_eq!( + decode_entries(records.clone()), + raw_entries(&records, Some(DecodeError::Malformed)) + ); + } + + #[rstest] + #[case::seconds("2026-08-31T12:34:56Z", 0)] + #[case::milliseconds("2026-08-31T12:34:56.789Z", 789_000_000)] + #[case::millisecond_whole("2026-08-31T12:34:56.000Z", 0)] + #[case::microseconds("2026-08-31T12:34:56.789123Z", 789_123_000)] + #[case::microsecond_trailing_zeros( + "2026-08-31T12:34:56.789000Z", + 789_000_000 + )] + #[case::nanoseconds("2026-08-31T12:34:56.789123456Z", 789_123_456)] + #[case::two_fraction_digits("2026-08-31T12:34:56.78Z", 780_000_000)] + #[case::four_fraction_digits("2026-08-31T12:34:56.7890Z", 789_000_000)] + #[case::numeric_offset("2026-08-31T14:34:56.789+02:00", 789_000_000)] + #[case::lowercase("2026-08-31t12:34:56.789z", 789_000_000)] + #[case::subnanoseconds("2026-08-31T12:34:56.7891234567Z", 789_123_456)] + fn native_format_accepts_rfc3339_timestamps( + #[case] timestamp: &str, + #[case] expected_nanos: u32, + ) { + let key = with_field(&key(0), 6, timestamp); + let diagnostic = + only_diagnostic(decode_entries(vec![(key, "value".into())])); + assert_eq!( + diagnostic + .key() + .timestamp + .to_rfc3339_opts(SecondsFormat::Nanos, true), + format!("2026-08-31T12:34:56.{expected_nanos:09}Z") + ); + } + + #[rstest] + #[case([1, 0, 2])] + #[case([2, 1, 0])] + fn chunk_permutations_decode_in_index_order(#[case] order: [u64; 3]) { + let records = order + .into_iter() + .map(|index| (key(index), ["a", "é", "c"][index as usize].into())) + .collect(); + let diagnostic = only_diagnostic(decode_entries(records)); + assert_eq!( + diagnostic.payload(), + &DiagnosticPayload::Text("aéc".into()) + ); + } + + #[test] + fn logical_groups_use_first_seen_order_not_timestamps() { + let later = with_field(&key(0), 6, "2026-08-31T12:35:00.000Z"); + let earlier = with_field(&key(0), 4, "earlier"); + let records = vec![ + (with_field(&later, 10, "1"), "b".into()), + ("raw".into(), "untouched".into()), + (earlier, "early".into()), + (later, "a".into()), + ]; + let entries = decode_entries(records); + assert!(matches!( + entries.as_slice(), + [Entry::Diagnostic(first), Entry::Raw(raw), Entry::Diagnostic(last)] + if first.payload() == &DiagnosticPayload::Text("ab".into()) + && raw.key == "raw" + && first.key().timestamp > last.key().timestamp + )); + } + + #[rstest] + #[case(vec![1])] + #[case(vec![0, 2])] + #[case(vec![2, 0])] + #[case(vec![u64::MAX])] + #[case(vec![0, u64::MAX])] + fn missing_chunks_preserve_all_members(#[case] indices: Vec) { + let records: Vec<_> = indices + .into_iter() + .map(|index| (key(index), format!("chunk {index}"))) + .collect(); + assert_eq!( + decode_entries(records.clone()), + raw_entries(&records, Some(DecodeError::IncompleteGroup)) + ); + } + + #[rstest] + #[case(vec![0, 0])] + #[case(vec![1, 1])] + #[case(vec![0, 1, 1])] + #[case(vec![2, 0, 2])] + fn duplicates_take_precedence_over_gaps(#[case] indices: Vec) { + let records: Vec<_> = indices + .into_iter() + .enumerate() + .map(|(position, index)| { + (key(index), format!("original {position}")) + }) + .collect(); + assert_eq!( + decode_entries(records.clone()), + raw_entries(&records, Some(DecodeError::DuplicateChunk)) + ); + } + + #[test] + fn duplicate_index_spellings_preserve_the_exact_keys() { + let records = vec![ + (key(0), "first".into()), + (with_field(&key(0), 10, "00"), "second".into()), + ]; + assert_eq!( + decode_entries(records.clone()), + raw_entries(&records, Some(DecodeError::DuplicateChunk)) + ); + } + + #[test] + fn failed_groups_preserve_physical_positions_among_other_entries() { + let records = vec![ + (key(2), "first seen".into()), + ("unrelated".into(), "unchanged".into()), + (key(0), "zero".into()), + (with_field(&key(0), 4, "valid"), "good".into()), + (key(2), "duplicate".into()), + ]; + let entries = decode_entries(records.clone()); + assert_eq!(entries.len(), records.len()); + for position in [0, 2, 4] { + let expected = raw_entries( + &records[position..=position], + Some(DecodeError::DuplicateChunk), + ); + assert_eq!(entries[position], expected[0]); + } + assert_eq!(entries[1], raw_entries(&records[1..2], None)[0]); + assert!( + matches!(&entries[3], Entry::Diagnostic(d) if d.key().name == "valid") + ); + } + + #[rstest] + #[case(1, "other-agent")] + #[case(2, "00000000-0000-0000-0000-000000000000")] + #[case(3, "start")] + #[case(4, "other-name")] + #[case(5, "00000000-0000-0000-0000-000000000000")] + #[case(6, "2026-08-31T12:34:56.790Z")] + #[case(7, "gz+b64")] + #[case(8, "success")] + #[case(9, "52")] + fn grouping_does_not_combine_different_metadata( + #[case] field: usize, + #[case] value: &str, + ) { + let detached = (with_field(&key(1), field, value), "detached".into()); + let entries = decode_entries(vec![ + (key(0), "a".into()), + detached.clone(), + (key(1), "b".into()), + ]); + assert_eq!(entries.len(), 2); + assert!(matches!( + &entries[0], + Entry::Diagnostic(d) if d.payload() == &DiagnosticPayload::Text("ab".into()) + )); + assert_eq!( + entries[1], + raw_entries(&[detached], Some(DecodeError::IncompleteGroup))[0] + ); + } + + #[rstest] + #[case(2, VM_ID, "3F2504E0-4F89-41D3-9A0C-0305E82C3301")] + #[case(5, EVENT_ID, "8f3e9c4a1b2c4d5e9f01234567890abc")] + #[case(9, "52", "052")] + fn equivalent_metadata_spellings_remain_separate_groups( + #[case] field: usize, + #[case] first: &str, + #[case] second: &str, + ) { + let records = vec![ + (with_field(&key(0), field, first), "a".into()), + (with_field(&key(0), field, second), "b".into()), + (with_field(&key(1), field, first), "c".into()), + (with_field(&key(1), field, second), "d".into()), + ]; + let entries = decode_entries(records); + assert!(matches!( + entries.as_slice(), + [Entry::Diagnostic(first), Entry::Diagnostic(second)] + if first.payload() == &DiagnosticPayload::Text("ac".into()) + && second.payload() == &DiagnosticPayload::Text("bd".into()) + )); + } + + #[test] + fn versions_are_never_grouped_together() { + let future = (with_field(&key(1), 0, "DIAG_V2"), "future chunk".into()); + let entries = decode_entries(vec![ + (key(0), "a".into()), + future.clone(), + (key(1), "b".into()), + ]); + assert_eq!(entries.len(), 2); + assert!(matches!( + &entries[0], + Entry::Diagnostic(d) if d.payload() == &DiagnosticPayload::Text("ab".into()) + )); + assert_eq!( + entries[1], + raw_entries(&[future], Some(DecodeError::UnsupportedVersion))[0] + ); + } + + #[rstest] + #[case("base64")] + #[case("zstd+b64")] + #[case("GZ+B64")] + fn unknown_encodings_preserve_every_chunk(#[case] encoding: &str) { + let records = vec![ + (with_field(&key(1), 7, encoding), "b".into()), + (with_field(&key(0), 7, encoding), "a".into()), + ]; + assert_eq!( + decode_entries(records.clone()), + raw_entries(&records, Some(DecodeError::Undecodable)) + ); + } + + #[test] + fn compressed_payload_is_joined_before_decoding() { + let records = vec![ + (with_field(&key(1), 7, "gz+b64"), GZIP_HELLO[5..].into()), + (with_field(&key(0), 7, "gz+b64"), GZIP_HELLO[..5].into()), + ]; + let diagnostic = only_diagnostic(decode_entries(records)); + assert_eq!(diagnostic.key().encoding, Some(Encoding::GzB64)); + assert_eq!( + diagnostic.payload(), + &DiagnosticPayload::Bytes(b"hello".to_vec()) + ); + } + + #[test] + fn corrupt_gzip_preserves_every_original_chunk() { + let mut gzip = STANDARD.decode(GZIP_HELLO).unwrap(); + let checksum_offset = gzip.len() - 8; + gzip[checksum_offset] ^= 1; + let value = STANDARD.encode(gzip); + let records = vec![ + (with_field(&key(1), 7, "gz+b64"), value[5..].into()), + (with_field(&key(0), 7, "gz+b64"), value[..5].into()), + ]; + assert_eq!( + decode_entries(records.clone()), + raw_entries(&records, Some(DecodeError::Undecodable)) + ); + } + + #[test] + fn visible_gaps_are_reported_before_payload_errors() { + let records = + vec![(with_field(&key(1), 7, "gz+b64"), "not base64".into())]; + assert_eq!( + decode_entries(records.clone()), + raw_entries(&records, Some(DecodeError::IncompleteGroup)) + ); + } + + #[test] + fn empty_text_payload_decodes() { + let diagnostic = + only_diagnostic(decode_entries(vec![(key(0), String::new())])); + assert_eq!( + diagnostic.payload(), + &DiagnosticPayload::Text(String::new()) + ); + } + + #[rstest] + #[case::agent(1, "a".repeat(64))] + #[case::name(4, "n".repeat(160))] + fn reader_accepts_fields_above_writer_budgets( + #[case] field: usize, + #[case] value: String, + ) { + let dir = TempDir::new().unwrap(); + let pool = + KvpPoolStore::new_in(KvpPool::Guest, dir.path(), PoolMode::Unsafe) + .unwrap(); + pool.append(&with_field(&key(0), field, &value), "payload") + .unwrap(); + let diagnostic = only_diagnostic( + DiagnosticReader::new(store(&dir)).entries().unwrap(), + ); + let actual = match field { + 1 => &diagnostic.key().agent, + _ => &diagnostic.key().name, + }; + assert_eq!(actual, &value); + } + + #[rstest] + #[case("0", Duration::ZERO)] + #[case("-0.0", Duration::ZERO)] + #[case("+1", Duration::from_secs(1))] + #[case("1.", Duration::from_secs(1))] + #[case(".5", Duration::from_millis(500))] + #[case("1.5", Duration::from_millis(1500))] + #[case("0.312000", Duration::from_millis(312))] + #[case("3.12e-1", Duration::from_millis(312))] + #[case("1.000000001", Duration::new(1, 1))] + #[case("0.9999999996", Duration::from_secs(1))] + #[case("1e-12", Duration::ZERO)] + #[case( + "9007199254740993.000000001", + Duration::new(9_007_199_254_740_993, 1) + )] + #[case("18446744073709551615.999999999", Duration::MAX)] + fn reader_accepts_duration_seconds( + #[case] seconds: &str, + #[case] expected: Duration, + ) { + let key = with_field(&key(0), 9, seconds); + let diagnostic = + only_diagnostic(decode_entries(vec![(key, "payload".into())])); + assert!(matches!(&diagnostic, Diagnostic::Event(event) + if event.duration == Some(expected))); + } + + #[test] + fn reader_accepts_values_above_safe_write_limit() { + let dir = TempDir::new().unwrap(); + let pool = + KvpPoolStore::new_in(KvpPool::Guest, dir.path(), PoolMode::Unsafe) + .unwrap(); + let value = "v".repeat(1500); + pool.append(&key(0), &value).unwrap(); + let diagnostic = only_diagnostic( + DiagnosticReader::new(store(&dir)).entries().unwrap(), + ); + assert_eq!(diagnostic.payload(), &DiagnosticPayload::Text(value)); + } + + #[test] + fn reader_does_not_limit_contiguous_group_size() { + let records = (0..1024) + .rev() + .map(|index| (key(index), "x".into())) + .collect(); + let diagnostic = only_diagnostic(decode_entries(records)); + assert_eq!( + diagnostic.payload(), + &DiagnosticPayload::Text("x".repeat(1024)) + ); + } +} diff --git a/libazureinit-kvp/src/diagnostics/writer.rs b/libazureinit-kvp/src/diagnostics/writer.rs new file mode 100644 index 00000000..8ac68b48 --- /dev/null +++ b/libazureinit-kvp/src/diagnostics/writer.rs @@ -0,0 +1,1219 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +use std::time::Duration; + +use chrono::{SecondsFormat, Utc}; +use uuid::Uuid; + +use super::diagnostic::{ + Diagnostic, DiagnosticEvent, DiagnosticFinish, DiagnosticKey, + DiagnosticPayload, DiagnosticStart, Encoding, Outcome, + DIAGNOSTIC_VERSION_ID, +}; +use super::encoding::encode_payload; +use super::MAX_CHUNK_BYTES; +use crate::{KvpError, KvpPoolStore}; + +const MAX_AGENT_BYTES: usize = 32; +const DEFAULT_MAX_NAME_BYTES: usize = 64; +const MAX_UUID_BYTES: usize = 36; +const MAX_TIMESTAMP_BYTES: usize = 30; +const MAX_KEY_BYTES: usize = 254; +const MAX_CHUNKS: usize = 1023; + +/// Fractional-second precision for emitted diagnostic timestamps. +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub enum TimestampPrecision { + /// Whole seconds, no fractional digits. + Seconds, + /// Millisecond precision (three fractional digits). The default. + #[default] + Millis, + /// Microsecond precision (six fractional digits). + Micros, + /// Nanosecond precision (nine fractional digits). + Nanos, +} + +impl TimestampPrecision { + fn seconds_format(self) -> SecondsFormat { + match self { + Self::Seconds => SecondsFormat::Secs, + Self::Millis => SecondsFormat::Millis, + Self::Micros => SecondsFormat::Micros, + Self::Nanos => SecondsFormat::Nanos, + } + } +} + +/// Fractional-second precision for emitted durations. +/// +/// Digits below the selected precision are discarded. +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub enum DurationPrecision { + /// Whole seconds, no fractional digits. + Seconds, + /// Milliseconds (three fractional digits). + Millis, + /// Microseconds (six fractional digits). The default. + #[default] + Micros, + /// Nanoseconds (nine fractional digits). + Nanos, +} + +impl DurationPrecision { + pub(crate) fn format(self, duration: Duration) -> String { + let digits = match self { + Self::Seconds => return duration.as_secs().to_string(), + Self::Millis => 3, + Self::Micros => 6, + Self::Nanos => 9, + }; + let fraction = duration.subsec_nanos() / 10u32.pow(9 - digits); + let width = digits as usize; + format!("{}.{fraction:0width$}", duration.as_secs()) + } +} + +/// Emits diagnostics for one reporting agent and VM. +/// +/// Use [`emit_event`](Self::emit_event) for a standalone observation. To record +/// an operation, call [`emit_start`](Self::emit_start) and +/// [`emit_finish`](Self::emit_finish) with the same event UUID. +/// +/// Timestamps and event UUIDs for standalone observations are generated +/// automatically. The caller supplies operation IDs, outcomes and durations. +/// Pass `None` for plain text or an [`Encoding`] for compressed text or bytes. +/// +/// Emission appends to the pool and splits large payloads automatically; it +/// never clears existing records. Invalid input writes nothing, while an I/O +/// error may leave part of a payload in the pool. See the +/// [diagnostics contract] for format details and limits. +/// +/// [diagnostics contract]: https://github.com/Azure/azure-init/blob/main/doc/diagnostics.md +/// +/// # Example +/// ```no_run +/// use std::time::Instant; +/// use libazureinit_kvp::{ +/// DiagnosticWriter, KvpPool, KvpPoolStore, Outcome, PoolMode, +/// }; +/// use uuid::Uuid; +/// +/// # fn main() -> Result<(), Box> { +/// let store = KvpPoolStore::new(KvpPool::Guest, PoolMode::Safe)?; +/// let writer = DiagnosticWriter::new( +/// store, "azure-init/0.1.1", "3f2504e0-4f89-41d3-9a0c-0305e82c3301", +/// )?; +/// let event_id = Uuid::new_v4().to_string(); +/// writer.emit_start(&event_id, "read:os-release", "reading OS information", None)?; +/// let started = Instant::now(); +/// let contents = std::fs::read_to_string("/etc/os-release")?; +/// writer.emit_finish( +/// &event_id, "read:os-release", format!("read {} bytes", contents.len()), +/// None, Outcome::Success, started.elapsed(), +/// )?; +/// # Ok(()) +/// # } +/// ``` +#[derive(Clone, Debug)] +pub struct DiagnosticWriter { + store: KvpPoolStore, + agent: String, + vm_id: String, + max_name_bytes: usize, + timestamp_precision: TimestampPrecision, + duration_precision: DurationPrecision, +} + +impl DiagnosticWriter { + /// Creates a writer with a fixed reporting agent and VM UUID. + /// + /// Versioned agents conventionally use `name/VERSION`. Invalid agent or + /// VM identifiers return [`KvpError`]; construction does not access the pool. + /// Timestamp precision defaults to milliseconds and duration precision to + /// microseconds. + pub fn new( + store: KvpPoolStore, + agent: impl Into, + vm_id: impl Into, + ) -> Result { + let agent = agent.into(); + let vm_id = vm_id.into(); + validate_field("agent", &agent, MAX_AGENT_BYTES)?; + validate_uuid("vm_id", &vm_id)?; + Ok(Self { + store, + agent, + vm_id, + max_name_bytes: DEFAULT_MAX_NAME_BYTES, + timestamp_precision: TimestampPrecision::default(), + duration_precision: DurationPrecision::default(), + }) + } + + /// Sets the maximum name length in UTF-8 bytes (default: 64). + /// + /// Names over the limit are rejected, never truncated. The complete key + /// must still fit within 254 bytes, including its chunk index. + pub fn with_max_name_bytes(mut self, max: usize) -> Self { + self.max_name_bytes = max; + self + } + + /// Overrides the default millisecond precision for emitted timestamps. + pub fn with_timestamp_precision( + mut self, + precision: TimestampPrecision, + ) -> Self { + self.timestamp_precision = precision; + self + } + + /// Sets duration precision independently of timestamp precision. + pub fn with_duration_precision( + mut self, + precision: DurationPrecision, + ) -> Self { + self.duration_precision = precision; + self + } + + /// Records the beginning of an operation. + /// + /// Pass the same event UUID and name to + /// [`emit_finish`](Self::emit_finish). + pub fn emit_start( + &self, + event_id: &str, + name: &str, + payload: impl Into, + encoding: Option, + ) -> Result<(), KvpError> { + self.emit(Diagnostic::Start(DiagnosticStart { + key: self.key(event_id, name, encoding), + payload: payload.into(), + })) + } + + /// Records an operation's outcome and caller-measured elapsed time. + /// + /// Use the event UUID and name passed to [`emit_start`](Self::emit_start). + /// The writer does not verify that a start exists or calculate the duration. + pub fn emit_finish( + &self, + event_id: &str, + name: &str, + payload: impl Into, + encoding: Option, + result: Outcome, + duration: Duration, + ) -> Result<(), KvpError> { + self.emit(Diagnostic::Finish(DiagnosticFinish { + key: self.key(event_id, name, encoding), + payload: payload.into(), + result, + duration, + })) + } + + /// Records a standalone observation with a generated event UUID. + /// + /// `result` and `duration` may be supplied independently. + pub fn emit_event( + &self, + name: &str, + payload: impl Into, + encoding: Option, + result: Option, + duration: Option, + ) -> Result<(), KvpError> { + self.emit(Diagnostic::Event(DiagnosticEvent { + key: self.key(&Uuid::new_v4().to_string(), name, encoding), + payload: payload.into(), + result, + duration, + })) + } + + fn key( + &self, + event_id: &str, + name: &str, + encoding: Option, + ) -> DiagnosticKey { + DiagnosticKey { + agent: self.agent.clone(), + vm_id: Some(self.vm_id.clone()), + name: name.to_owned(), + event_id: event_id.to_owned(), + timestamp: Utc::now(), + encoding, + } + } + + fn emit(&self, diagnostic: Diagnostic) -> Result<(), KvpError> { + self.store.append_multiple(prepare_records( + diagnostic, + self.timestamp_precision, + self.duration_precision, + self.max_name_bytes, + )?) + } +} + +fn prepare_records( + diagnostic: Diagnostic, + timestamp_precision: TimestampPrecision, + duration_precision: DurationPrecision, + max_name_bytes: usize, +) -> Result, KvpError> { + let kind = diagnostic.kind(); + let (key, payload, result, duration) = match diagnostic { + Diagnostic::Start(start) => (start.key, start.payload, None, None), + Diagnostic::Finish(finish) => ( + finish.key, + finish.payload, + Some(finish.result), + Some(finish.duration), + ), + Diagnostic::Event(event) => { + (event.key, event.payload, event.result, event.duration) + } + }; + + validate_field("agent", &key.agent, MAX_AGENT_BYTES)?; + let vm_id = key + .vm_id + .as_deref() + .ok_or(KvpError::EmptyEventField { field: "vm_id" })?; + validate_uuid("vm_id", vm_id)?; + validate_field("name", &key.name, max_name_bytes)?; + validate_uuid("event_id", &key.event_id)?; + + let duration = duration.map_or_else(String::new, |duration| { + duration_precision.format(duration) + }); + let timestamp = key + .timestamp + .to_rfc3339_opts(timestamp_precision.seconds_format(), true); + validate_field("timestamp", ×tamp, MAX_TIMESTAMP_BYTES)?; + + let value = encode_payload(payload, key.encoding.as_ref())?; + let encoding = key + .encoding + .as_ref() + .map_or_else(|| "none".to_owned(), ToString::to_string); + let result = result.map_or_else(String::new, |result| result.to_string()); + let base_key = format!( + "{DIAGNOSTIC_VERSION_ID}|{}|{vm_id}|{kind}|{}|{}|{timestamp}|{encoding}|{result}|{duration}", + key.agent, key.name, key.event_id, + ); + frame_records(&base_key, &value) +} + +fn validate_field( + field: &'static str, + value: &str, + max: usize, +) -> Result<(), KvpError> { + if value.is_empty() { + return Err(KvpError::EmptyEventField { field }); + } + if value.contains('|') { + return Err(KvpError::EventFieldContainsDelimiter { field }); + } + if value.contains('\0') { + return Err(KvpError::KeyContainsNull); + } + if value.len() > max { + return Err(KvpError::EventFieldTooLong { + field, + max, + actual: value.len(), + }); + } + Ok(()) +} + +fn validate_uuid(field: &'static str, value: &str) -> Result<(), KvpError> { + validate_field(field, value, MAX_UUID_BYTES)?; + Uuid::parse_str(value) + .map(|_| ()) + .map_err(|_| KvpError::InvalidUuid { field }) +} + +fn frame_records( + base_key: &str, + mut value: &str, +) -> Result, KvpError> { + let mut records = Vec::new(); + loop { + if records.len() == MAX_CHUNKS { + return Err(KvpError::TooManyChunks { max: MAX_CHUNKS }); + } + let key = format!("{base_key}|{}", records.len()); + if key.len() > MAX_KEY_BYTES { + return Err(KvpError::KeyTooLarge { + max: MAX_KEY_BYTES, + actual: key.len(), + }); + } + + let mut end = value.len().min(MAX_CHUNK_BYTES); + while !value.is_char_boundary(end) { + end -= 1; + } + let (chunk, remaining) = value.split_at(end); + records.push((key, chunk.to_owned())); + if remaining.is_empty() { + return Ok(records); + } + value = remaining; + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::fs; + use std::io; + use std::path::Path; + use std::sync::atomic::{AtomicUsize, Ordering}; + use std::sync::Arc; + + use chrono::DateTime; + use rstest::rstest; + use tempfile::TempDir; + + use super::super::encoding::decode_payload; + use crate::store::{Handle, OsSysOps, StatInfo, SysOps}; + use crate::{DiagnosticReader, Entry, KvpPool, PoolMode}; + + const AGENT: &str = "azure-init/0.1.1"; + const VM_ID: &str = "3f2504e0-4f89-41d3-9a0c-0305e82c3301"; + const EVENT_ID: &str = "8f3e9c4a-1b2c-4d5e-9f01-234567890abc"; + const TIMESTAMP: &str = "2026-08-31T12:34:56.789Z"; + + fn store(dir: &TempDir, mode: PoolMode) -> KvpPoolStore { + KvpPoolStore::new_in(KvpPool::Guest, dir.path(), mode).unwrap() + } + + fn key() -> DiagnosticKey { + DiagnosticKey { + agent: AGENT.into(), + vm_id: Some(VM_ID.into()), + name: "provision:run".into(), + event_id: EVENT_ID.into(), + timestamp: DateTime::parse_from_rfc3339(TIMESTAMP) + .unwrap() + .with_timezone(&Utc), + encoding: None, + } + } + + fn event(key: DiagnosticKey, payload: DiagnosticPayload) -> Diagnostic { + Diagnostic::Event(DiagnosticEvent { + key, + payload, + result: None, + duration: None, + }) + } + + #[derive(Debug, Default)] + struct WriterOps { + os: OsSysOps, + calls: AtomicUsize, + } + + impl SysOps for WriterOps { + fn open_read(&self, _: &Path) -> io::Result> { + self.calls.fetch_add(1, Ordering::SeqCst); + Err(io::ErrorKind::Unsupported.into()) + } + + fn open_read_write(&self, _: &Path) -> io::Result> { + self.calls.fetch_add(1, Ordering::SeqCst); + Err(io::ErrorKind::Unsupported.into()) + } + + fn open_read_write_create( + &self, + path: &Path, + ) -> io::Result> { + self.calls.fetch_add(1, Ordering::SeqCst); + self.os.open_read_write_create(path) + } + + fn path_metadata(&self, _: &Path) -> io::Result { + self.calls.fetch_add(1, Ordering::SeqCst); + Err(io::ErrorKind::Unsupported.into()) + } + + fn boot_time(&self) -> io::Result { + self.calls.fetch_add(1, Ordering::SeqCst); + Err(io::ErrorKind::Unsupported.into()) + } + } + + fn observed_writer(dir: &TempDir) -> (DiagnosticWriter, Arc) { + let ops = Arc::new(WriterOps::default()); + let store = KvpPoolStore::with_ops( + KvpPool::Guest, + dir.path(), + PoolMode::Safe, + ops.clone(), + ) + .unwrap(); + (DiagnosticWriter::new(store, AGENT, VM_ID).unwrap(), ops) + } + + #[test] + fn writer_ops_rejects_non_append_operations() { + let dir = TempDir::new().unwrap(); + let pool = store(&dir, PoolMode::Safe); + let ops = WriterOps::default(); + assert_eq!( + [ + ops.open_read(pool.path()).unwrap_err().kind(), + ops.open_read_write(pool.path()).unwrap_err().kind(), + ops.path_metadata(pool.path()).unwrap_err().kind(), + ops.boot_time().unwrap_err().kind(), + ], + [io::ErrorKind::Unsupported; 4] + ); + assert_eq!(ops.calls.load(Ordering::SeqCst), 4); + assert!(!pool.path().exists()); + } + + fn assert_rejected_without_writes( + operation: impl FnOnce(&DiagnosticWriter) -> Result<(), KvpError>, + ) -> KvpError { + let dir = TempDir::new().unwrap(); + let pool = store(&dir, PoolMode::Safe); + pool.append("existing", "value").unwrap(); + let before = fs::read(pool.path()).unwrap(); + let (writer, ops) = observed_writer(&dir); + let error = operation(&writer).unwrap_err(); + assert_eq!(ops.calls.load(Ordering::SeqCst), 0); + assert_eq!(fs::read(pool.path()).unwrap(), before); + error + } + + #[test] + fn constructor_does_no_io_and_emission_opens_once() { + let dir = TempDir::new().unwrap(); + let pool = store(&dir, PoolMode::Safe); + let (writer, ops) = observed_writer(&dir); + assert_eq!(ops.calls.load(Ordering::SeqCst), 0); + assert!(!pool.path().exists()); + + let payload = "x".repeat(MAX_CHUNK_BYTES * 3 + 1); + writer + .emit_event("test", payload, None, None, None) + .unwrap(); + assert_eq!(ops.calls.load(Ordering::SeqCst), 1); + assert_eq!(pool.dump().unwrap().len(), 4); + } + + #[test] + fn constructor_rejects_empty_agent_without_io() { + let dir = TempDir::new().unwrap(); + let (writer, ops) = observed_writer(&dir); + assert!(matches!( + DiagnosticWriter::new(writer.store, "", VM_ID), + Err(KvpError::EmptyEventField { field: "agent" }) + )); + assert_eq!(ops.calls.load(Ordering::SeqCst), 0); + } + + #[test] + fn constructor_rejects_invalid_vm_id_without_io() { + let dir = TempDir::new().unwrap(); + let (writer, ops) = observed_writer(&dir); + assert!(matches!( + DiagnosticWriter::new(writer.store, AGENT, "vm-abc"), + Err(KvpError::InvalidUuid { field: "vm_id" }) + )); + assert_eq!(ops.calls.load(Ordering::SeqCst), 0); + } + + #[rstest] + #[case("invalid")] + #[case("3f2504e0-4f89-41d3-9a0c-0305e82c330z")] + fn uuid_validation_rejects_invalid_syntax(#[case] value: &str) { + assert!(matches!( + validate_uuid("event_id", value), + Err(KvpError::InvalidUuid { field: "event_id" }) + )); + } + + #[test] + fn uuid_validation_preserves_field_errors() { + assert!(matches!( + validate_uuid("event_id", ""), + Err(KvpError::EmptyEventField { field: "event_id" }) + )); + assert!(matches!( + validate_uuid("event_id", "bad|id"), + Err(KvpError::EventFieldContainsDelimiter { field: "event_id" }) + )); + assert!(matches!( + validate_uuid("event_id", "bad\0id"), + Err(KvpError::KeyContainsNull) + )); + } + + #[test] + fn uuid_validation_rejects_oversized_representations() { + assert!(matches!( + validate_uuid("vm_id", &format!("{{{VM_ID}}}")), + Err(KvpError::EventFieldTooLong { + field: "vm_id", + max: MAX_UUID_BYTES, + actual: 38, + }) + )); + } + + #[rstest] + #[case(VM_ID)] + #[case("3F2504E0-4F89-41D3-9A0C-0305E82C3301")] + #[case("3f2504e04f8941d39a0c0305e82c3301")] + #[case("00000000-0000-0000-0000-000000000000")] + fn valid_uuid_spellings_are_preserved(#[case] id: &str) { + let dir = TempDir::new().unwrap(); + let pool = store(&dir, PoolMode::Safe); + let writer = DiagnosticWriter::new(pool.clone(), AGENT, id).unwrap(); + writer.emit_start(id, "test", "starting", None).unwrap(); + let records = pool.dump().unwrap(); + let fields: Vec<_> = records[0].0.split('|').collect(); + assert_eq!(fields[2], id); + assert_eq!(fields[5], id); + } + + #[rstest] + #[case("agent", "a".repeat(32), 32, None)] + #[case("agent", "é".repeat(16), 32, None)] + #[case("name", "n".repeat(64), 64, None)] + #[case("name", "é".repeat(32), 64, None)] + #[case("name", "n".repeat(32), 32, Some(32))] + #[case("name", "n".repeat(96), 96, Some(96))] + fn freeform_caps_count_bytes_without_truncation( + #[case] field: &'static str, + #[case] value: String, + #[case] max: usize, + #[case] max_name_bytes: Option, + #[values(PoolMode::Safe, PoolMode::Unsafe)] mode: PoolMode, + ) { + let dir = TempDir::new().unwrap(); + let pool = store(&dir, mode); + let (agent, name) = match field { + "agent" => (value.as_str(), "test"), + _ => (AGENT, value.as_str()), + }; + let writer = DiagnosticWriter::new(pool.clone(), agent, VM_ID).unwrap(); + let writer = match max_name_bytes { + Some(max) => writer.with_max_name_bytes(max), + None => writer, + }; + writer.emit_event(name, "ok", None, None, None).unwrap(); + let records = pool.dump().unwrap(); + let fields: Vec<_> = records[0].0.split('|').collect(); + assert_eq!(fields[1], agent); + assert_eq!(fields[4], name); + + let before = fs::read(pool.path()).unwrap(); + let oversized = format!("{value}x"); + let error = match field { + "agent" => DiagnosticWriter::new(pool.clone(), oversized, VM_ID) + .unwrap_err(), + _ => writer + .emit_event(&oversized, "bad", None, None, None) + .unwrap_err(), + }; + assert!(matches!( + error, + KvpError::EventFieldTooLong { field: actual_field, max: cap, actual } + if actual_field == field && cap == max && actual == max + 1 + )); + assert_eq!(fs::read(pool.path()).unwrap(), before); + } + + #[rstest] + #[case("", "empty")] + #[case("bad|field", "delimiter")] + #[case("bad\0field", "null")] + fn field_validation_reports_the_reason( + #[case] value: &str, + #[case] reason: &str, + ) { + let error = + validate_field("name", value, DEFAULT_MAX_NAME_BYTES).unwrap_err(); + match reason { + "empty" => { + assert!(matches!( + error, + KvpError::EmptyEventField { field: "name" } + )) + } + "delimiter" => assert!(matches!( + error, + KvpError::EventFieldContainsDelimiter { field: "name" } + )), + _ => assert!(matches!(error, KvpError::KeyContainsNull)), + } + } + + #[test] + fn exact_start_key_matches_diag_layout() { + let records = prepare_records( + Diagnostic::Start(DiagnosticStart { + key: key(), + payload: "starting".into(), + }), + TimestampPrecision::Millis, + DurationPrecision::default(), + DEFAULT_MAX_NAME_BYTES, + ) + .unwrap(); + assert_eq!( + records, + vec![( + format!( + "DIAG|{AGENT}|{VM_ID}|start|provision:run|{EVENT_ID}|{TIMESTAMP}|none|||0" + ), + "starting".to_owned(), + )] + ); + } + + #[rstest] + #[case::success(Outcome::Success, "success", 312_000, "0.312000")] + #[case::failure(Outcome::Failure, "fail", 0, "0.000000")] + fn exact_finish_key_matches_diag_layout( + #[case] result: Outcome, + #[case] token: &str, + #[case] duration_us: u64, + #[case] seconds: &str, + ) { + let records = prepare_records( + Diagnostic::Finish(DiagnosticFinish { + key: key(), + payload: "finished".into(), + result, + duration: Duration::from_micros(duration_us), + }), + TimestampPrecision::Millis, + DurationPrecision::default(), + DEFAULT_MAX_NAME_BYTES, + ) + .unwrap(); + assert_eq!( + records[0].0, + format!( + "DIAG|{AGENT}|{VM_ID}|finish|provision:run|{EVENT_ID}|{TIMESTAMP}|none|{token}|{seconds}|0" + ) + ); + } + + #[rstest] + #[case::neither(None, None, "")] + #[case::result_only(Some(Outcome::Success), None, "")] + #[case::zero_duration(None, Some(0), "0.000000")] + #[case::both(Some(Outcome::Failure), Some(52), "0.000052")] + fn event_optional_fields_are_independent( + #[case] result: Option, + #[case] duration_us: Option, + #[case] seconds: &str, + ) { + let dir = TempDir::new().unwrap(); + let pool = store(&dir, PoolMode::Safe); + let writer = DiagnosticWriter::new(pool.clone(), AGENT, VM_ID).unwrap(); + writer + .emit_event( + "test", + "ok", + None, + result, + duration_us.map(Duration::from_micros), + ) + .unwrap(); + let records = pool.dump().unwrap(); + let fields: Vec<_> = records[0].0.split('|').collect(); + assert_eq!(fields.len(), 11); + assert_eq!(fields[3], "event"); + assert_eq!( + fields[8], + result.map_or_else(String::new, |v| v.to_string()) + ); + assert_eq!(fields[9], seconds); + } + + #[test] + fn span_endpoints_reuse_the_supplied_event_id() { + let dir = TempDir::new().unwrap(); + let pool = store(&dir, PoolMode::Safe); + let writer = DiagnosticWriter::new(pool.clone(), AGENT, VM_ID).unwrap(); + writer + .emit_start(EVENT_ID, "test", "starting", None) + .unwrap(); + writer + .emit_finish( + EVENT_ID, + "test", + "finished", + None, + Outcome::Failure, + Duration::from_micros(17), + ) + .unwrap(); + let records = pool.dump().unwrap(); + assert_eq!(records.len(), 2); + for (record, kind) in records.iter().zip(["start", "finish"]) { + let fields: Vec<_> = record.0.split('|').collect(); + assert_eq!(fields[3], kind); + assert_eq!(fields[5], EVENT_ID); + } + } + + #[test] + fn events_receive_distinct_v4_ids() { + let dir = TempDir::new().unwrap(); + let pool = store(&dir, PoolMode::Safe); + let writer = DiagnosticWriter::new(pool.clone(), AGENT, VM_ID).unwrap(); + for _ in 0..2 { + writer + .emit_event("test", "message", None, None, None) + .unwrap(); + } + let ids: Vec<_> = pool + .dump() + .unwrap() + .iter() + .map(|(key, _)| { + Uuid::parse_str(key.split('|').nth(5).unwrap()).unwrap() + }) + .collect(); + assert_eq!(ids.len(), 2); + assert_ne!(ids[0], ids[1]); + for id in ids { + assert_eq!(id.get_version_num(), 4); + } + } + + #[test] + fn emission_uses_a_current_utc_millisecond_timestamp() { + let dir = TempDir::new().unwrap(); + let pool = store(&dir, PoolMode::Safe); + let writer = DiagnosticWriter::new(pool.clone(), AGENT, VM_ID).unwrap(); + let before = Utc::now().timestamp_millis(); + writer + .emit_event("test", "message", None, None, None) + .unwrap(); + let after = Utc::now().timestamp_millis(); + let records = pool.dump().unwrap(); + assert_eq!(records.len(), 1); + let timestamp = records[0].0.split('|').nth(6).unwrap(); + assert_eq!(timestamp.len(), 24); + assert!(timestamp.ends_with('Z')); + let parsed = DateTime::parse_from_rfc3339(timestamp).unwrap(); + assert!((before..=after).contains(&parsed.timestamp_millis())); + } + + #[rstest] + #[case(TimestampPrecision::Seconds, 20)] + #[case(TimestampPrecision::Millis, 24)] + #[case(TimestampPrecision::Micros, 27)] + #[case(TimestampPrecision::Nanos, 30)] + fn timestamp_precision_controls_emitted_width( + #[case] precision: TimestampPrecision, + #[case] expected_len: usize, + ) { + let dir = TempDir::new().unwrap(); + let pool = store(&dir, PoolMode::Safe); + let writer = DiagnosticWriter::new(pool.clone(), AGENT, VM_ID) + .unwrap() + .with_timestamp_precision(precision); + writer.emit_event("test", "ok", None, None, None).unwrap(); + let records = pool.dump().unwrap(); + let timestamp = records[0].0.split('|').nth(6).unwrap(); + assert_eq!(timestamp.len(), expected_len); + assert!(timestamp.ends_with('Z')); + } + + #[rstest] + #[case(DurationPrecision::Seconds, "1", Duration::from_secs(1))] + #[case(DurationPrecision::Millis, "1.123", Duration::from_millis(1123))] + #[case( + DurationPrecision::Micros, + "1.123456", + Duration::from_micros(1_123_456) + )] + #[case( + DurationPrecision::Nanos, + "1.123456789", + Duration::new(1, 123_456_789) + )] + fn duration_precision_controls_emitted_seconds( + #[case] precision: DurationPrecision, + #[case] seconds: &str, + #[case] expected: Duration, + ) { + let dir = TempDir::new().unwrap(); + let pool = store(&dir, PoolMode::Safe); + let writer = DiagnosticWriter::new(pool.clone(), AGENT, VM_ID) + .unwrap() + .with_duration_precision(precision); + writer + .emit_event( + "timed", + "ok", + None, + None, + Some(Duration::new(1, 123_456_789)), + ) + .unwrap(); + let records = pool.dump().unwrap(); + let fields: Vec<_> = records[0].0.split('|').collect(); + assert_eq!(fields[9], seconds); + assert_eq!(fields[6].len(), 24); + let entries = DiagnosticReader::new(pool).entries().unwrap(); + assert!( + matches!(&entries[0], Entry::Diagnostic(Diagnostic::Event(event)) + if event.duration == Some(expected)) + ); + } + + #[rstest] + #[case(String::new(), vec![0])] + #[case("x".repeat(MAX_CHUNK_BYTES), vec![MAX_CHUNK_BYTES])] + #[case("x".repeat(MAX_CHUNK_BYTES + 1), vec![MAX_CHUNK_BYTES, 1])] + #[case(format!("{}é", "x".repeat(1021)), vec![1021, 2])] + #[case(format!("{}€", "x".repeat(1021)), vec![1021, 3])] + #[case(format!("{}😀", "x".repeat(1021)), vec![1021, 4])] + #[case("€".repeat(1000), vec![1020, 1020, 960])] + fn framing_preserves_utf8_and_empty_values( + #[case] value: String, + #[case] lengths: Vec, + ) { + let records = frame_records("base", &value).unwrap(); + assert_eq!( + records + .iter() + .map(|(_, value)| value.len()) + .collect::>(), + lengths + ); + assert_eq!( + records + .iter() + .map(|(_, value)| value.as_str()) + .collect::(), + value + ); + for (index, (key, value)) in records.iter().enumerate() { + assert_eq!(key, &format!("base|{index}")); + assert!(value.len() <= MAX_CHUNK_BYTES); + } + } + + #[rstest] + fn compression_precedes_safe_framing_in_both_modes( + #[values(PoolMode::Safe, PoolMode::Unsafe)] mode: PoolMode, + #[values(Encoding::GzB64, Encoding::ZlibB64)] encoding: Encoding, + ) { + let dir = TempDir::new().unwrap(); + let pool = store(&dir, mode); + let writer = DiagnosticWriter::new(pool.clone(), AGENT, VM_ID).unwrap(); + let bytes: Vec<_> = (0..4096u32).flat_map(u32::to_le_bytes).collect(); + writer + .emit_event( + "provision:run", + bytes.clone(), + Some(encoding.clone()), + None, + None, + ) + .unwrap(); + let records = pool.dump().unwrap(); + assert!(records.len() > 1); + let base = records[0].0.rsplit_once('|').unwrap().0; + for (index, (key, value)) in records.iter().enumerate() { + assert_eq!(key, &format!("{base}|{index}")); + let fields: Vec<_> = key.split('|').collect(); + assert_eq!(fields[3], "event"); + assert_eq!(fields[7], encoding.to_string()); + assert!(key.len() <= MAX_KEY_BYTES); + assert!(value.len() <= MAX_CHUNK_BYTES); + } + let value: String = records.iter().map(|(_, v)| v.as_str()).collect(); + assert_eq!( + decode_payload(value.as_bytes(), Some(&encoding)).unwrap(), + DiagnosticPayload::Bytes(bytes) + ); + } + + #[test] + fn maximum_chunk_count_is_accepted() { + let dir = TempDir::new().unwrap(); + let pool = store(&dir, PoolMode::Safe); + let writer = DiagnosticWriter::new(pool.clone(), AGENT, VM_ID).unwrap(); + let payload = "x".repeat(MAX_CHUNK_BYTES * MAX_CHUNKS); + writer + .emit_event("test", payload, None, None, None) + .unwrap(); + let records = pool.dump().unwrap(); + assert_eq!(records.len(), MAX_CHUNKS); + assert!(records.last().unwrap().0.ends_with("|1022")); + } + + #[rstest] + #[case("x".repeat(MAX_CHUNK_BYTES * MAX_CHUNKS + 1))] + #[case("€".repeat((MAX_CHUNK_BYTES / 3) * MAX_CHUNKS + 1))] + fn excess_chunks_do_not_open_or_modify_the_pool(#[case] payload: String) { + let dir = TempDir::new().unwrap(); + let pool = store(&dir, PoolMode::Safe); + pool.append("existing", "value").unwrap(); + let before = fs::read(pool.path()).unwrap(); + let (writer, ops) = observed_writer(&dir); + assert!(matches!( + writer.emit_event("test", payload, None, None, None), + Err(KvpError::TooManyChunks { max: MAX_CHUNKS }) + )); + assert_eq!(ops.calls.load(Ordering::SeqCst), 0); + assert_eq!(fs::read(pool.path()).unwrap(), before); + } + + #[test] + fn chunk_limit_applies_after_compression() { + let dir = TempDir::new().unwrap(); + let pool = store(&dir, PoolMode::Safe); + let writer = DiagnosticWriter::new(pool.clone(), AGENT, VM_ID).unwrap(); + let bytes = vec![0; MAX_CHUNK_BYTES * MAX_CHUNKS + 1]; + writer + .emit_event( + "test", + bytes.clone(), + Some(Encoding::GzB64), + None, + None, + ) + .unwrap(); + let records = pool.dump().unwrap(); + assert!(records.len() < MAX_CHUNKS); + let value: String = records.iter().map(|(_, v)| v.as_str()).collect(); + assert_eq!( + decode_payload(value.as_bytes(), Some(&Encoding::GzB64)).unwrap(), + DiagnosticPayload::Bytes(bytes) + ); + } + + #[test] + fn key_limit_includes_every_chunk_suffix() { + let base = "x".repeat(MAX_KEY_BYTES - 2); + let records = frame_records(&base, "").unwrap(); + assert_eq!(records[0].0.len(), MAX_KEY_BYTES); + let payload = "v".repeat(MAX_CHUNK_BYTES * 11); + assert!(matches!( + frame_records(&base, &payload), + Err(KvpError::KeyTooLarge { + max: MAX_KEY_BYTES, + actual: 255 + }) + )); + } + + #[test] + fn default_name_limit_fits_with_four_digit_index() { + let diagnostic = Diagnostic::Finish(DiagnosticFinish { + key: DiagnosticKey { + agent: "a".repeat(MAX_AGENT_BYTES), + name: "n".repeat(DEFAULT_MAX_NAME_BYTES), + encoding: Some(Encoding::ZlibB64), + ..key() + }, + payload: "test".into(), + result: Outcome::Success, + duration: Duration::from_millis(312), + }); + let records = prepare_records( + diagnostic, + TimestampPrecision::Nanos, + DurationPrecision::Nanos, + DEFAULT_MAX_NAME_BYTES, + ) + .unwrap(); + let base = records[0].0.rsplit_once('|').unwrap().0; + let longest = format!("{base}|1022"); + assert_eq!(longest.len(), 248); + assert!(longest.len() <= MAX_KEY_BYTES); + } + + #[test] + fn name_allowance_does_not_override_full_key_limit() { + let error = assert_rejected_without_writes(|writer| { + DiagnosticWriter::new( + writer.store.clone(), + "a".repeat(MAX_AGENT_BYTES), + VM_ID, + )? + .with_timestamp_precision(TimestampPrecision::Nanos) + .with_duration_precision(DurationPrecision::Nanos) + .emit_finish( + EVENT_ID, + &"n".repeat(DEFAULT_MAX_NAME_BYTES), + "message", + Some(Encoding::ZlibB64), + Outcome::Success, + Duration::MAX, + ) + }); + assert!(matches!( + error, + KvpError::KeyTooLarge { + max: MAX_KEY_BYTES, + actual: 264 + } + )); + } + + #[test] + fn start_rejects_invalid_event_id_before_io() { + let error = assert_rejected_without_writes(|writer| { + writer.emit_start("invalid", "test", "message", None) + }); + assert!(matches!(error, KvpError::InvalidUuid { field: "event_id" })); + } + + #[test] + fn finish_rejects_invalid_name_before_io() { + let error = assert_rejected_without_writes(|writer| { + writer.emit_finish( + EVENT_ID, + "bad|name", + "message", + None, + Outcome::Failure, + Duration::ZERO, + ) + }); + assert!(matches!( + error, + KvpError::EventFieldContainsDelimiter { field: "name" } + )); + } + + #[test] + fn event_rejects_empty_name_before_io() { + let error = assert_rejected_without_writes(|writer| { + writer.emit_event("", "message", None, None, None) + }); + assert!(matches!(error, KvpError::EmptyEventField { field: "name" })); + } + + #[test] + fn start_rejects_invalid_utf8_before_io() { + let error = assert_rejected_without_writes(|writer| { + writer.emit_start(EVENT_ID, "test", vec![0xff], None) + }); + assert!(matches!(error, KvpError::PayloadNotUtf8)); + } + + #[test] + fn finish_rejects_late_nul_before_writing_any_chunks() { + let payload = format!("{}\0", "x".repeat(MAX_CHUNK_BYTES)); + let error = assert_rejected_without_writes(|writer| { + writer.emit_finish( + EVENT_ID, + "test", + payload, + None, + Outcome::Failure, + Duration::ZERO, + ) + }); + assert!(matches!(error, KvpError::ValueContainsNull)); + } + + #[test] + fn event_rejects_unsupported_encoding_before_io() { + let error = assert_rejected_without_writes(|writer| { + writer.emit_event( + "test", + "message", + Some(Encoding::Other("zstd".into())), + None, + None, + ) + }); + assert!(matches!( + error, + KvpError::UnsupportedEncoding { token } if token == "zstd" + )); + } + + #[test] + fn writer_rejects_absent_vm_identity() { + let mut missing_vm = key(); + missing_vm.vm_id = None; + assert!(matches!( + prepare_records( + event(missing_vm, "bad".into()), + TimestampPrecision::Millis, + DurationPrecision::default(), + DEFAULT_MAX_NAME_BYTES, + ), + Err(KvpError::EmptyEventField { field: "vm_id" }) + )); + } + + #[test] + fn writer_rejects_expanded_timestamp_years() { + let mut expanded_year = key(); + expanded_year.timestamp = + DateTime::from_timestamp(253_402_300_800, 0).unwrap(); + assert!(matches!( + prepare_records( + event(expanded_year, "bad".into()), + TimestampPrecision::Nanos, + DurationPrecision::default(), + DEFAULT_MAX_NAME_BYTES, + ), + Err(KvpError::EventFieldTooLong { + field: "timestamp", + max: 30, + .. + }) + )); + } + + #[test] + fn appends_preserve_existing_records_and_duplicates() { + let dir = TempDir::new().unwrap(); + let pool = store(&dir, PoolMode::Safe); + pool.append("existing", "first").unwrap(); + pool.append("existing", "second").unwrap(); + let writer = DiagnosticWriter::new(pool.clone(), AGENT, VM_ID).unwrap(); + writer.emit_event("test", "new", None, None, None).unwrap(); + let records = pool.dump().unwrap(); + assert_eq!(records.len(), 3); + assert_eq!(records[0], ("existing".into(), "first".into())); + assert_eq!(records[1], ("existing".into(), "second".into())); + } + + #[test] + fn constructor_defers_storage_errors_until_emission() { + let dir = TempDir::new().unwrap(); + let pool = store(&dir, PoolMode::Safe); + fs::create_dir(pool.path()).unwrap(); + let writer = DiagnosticWriter::new(pool, AGENT, VM_ID).unwrap(); + assert!(matches!( + writer.emit_event("test", "value", None, None, None), + Err(KvpError::Io(_)) + )); + } +} diff --git a/libazureinit-kvp/src/error.rs b/libazureinit-kvp/src/error.rs index 02b9ec9c..e2ed734b 100644 --- a/libazureinit-kvp/src/error.rs +++ b/libazureinit-kvp/src/error.rs @@ -4,24 +4,71 @@ use std::fmt; use std::io; -/// Errors returned by [`KvpPoolStore`](crate::KvpPoolStore) operations. +/// Errors returned by KVP storage and telemetry writers. #[derive(Debug)] pub enum KvpError { /// The key was empty. EmptyKey, - /// An underlying I/O error. + /// A required diagnostic field was empty. + EmptyEventField { + /// Name of the rejected field. + field: &'static str, + }, + /// An I/O operation failed, or stored pool data was invalid. Io(io::Error), - /// The key contains a null byte, which is incompatible with the - /// on-disk format (null-padded fixed-width fields). + /// A diagnostic field contained the reserved `|` delimiter. + EventFieldContainsDelimiter { + /// Name of the rejected field. + field: &'static str, + }, + /// A diagnostic field exceeded its UTF-8 byte limit. + EventFieldTooLong { + /// Name of the rejected field. + field: &'static str, + /// Maximum allowed bytes. + max: usize, + /// Supplied bytes. + actual: usize, + }, + /// A VM or event identifier was not a valid UUID. + InvalidUuid { + /// Name of the rejected field. + field: &'static str, + }, + /// An encoded payload needed too many records. + TooManyChunks { + /// Maximum allowed records for one payload. + max: usize, + }, + /// A key or diagnostic key field contained a NUL byte. KeyContainsNull, /// The key exceeds the store's maximum key size. - KeyTooLarge { max: usize, actual: usize }, - /// The store already has the maximum allowed number of unique keys. - MaxUniqueKeysExceeded { max: usize }, + KeyTooLarge { + /// Maximum allowed UTF-8 bytes. + max: usize, + /// Supplied UTF-8 bytes. + actual: usize, + }, + /// An insert or replacement would exceed the unique-key limit. + MaxUniqueKeysExceeded { + /// Maximum allowed distinct keys. + max: usize, + }, + /// A byte payload could not be written as unencoded UTF-8 text. + PayloadNotUtf8, + /// The requested payload encoding is not supported. + UnsupportedEncoding { + /// Requested encoding name. + token: String, + }, /// The value exceeds the store's maximum value size. - ValueTooLarge { max: usize, actual: usize }, - /// The value contains a null byte, which is incompatible with the - /// null-padded KVP wire format. + ValueTooLarge { + /// Maximum allowed UTF-8 bytes. + max: usize, + /// Supplied UTF-8 bytes. + actual: usize, + }, + /// A stored value contained a NUL byte. ValueContainsNull, } @@ -29,6 +76,21 @@ 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::EmptyEventField { field } => { + write!(f, "event key field '{field}' must not be empty") + } + Self::EventFieldContainsDelimiter { field } => { + write!(f, "event key field '{field}' must not contain '|'") + } + Self::EventFieldTooLong { field, max, actual } => { + write!(f, "event key field '{field}' length ({actual}) exceeds maximum ({max})") + } + Self::InvalidUuid { field } => { + write!(f, "event key field '{field}' must be a UUID") + } + Self::TooManyChunks { max } => { + write!(f, "diagnostic chunk count exceeds maximum ({max})") + } Self::Io(e) => write!(f, "{e}"), Self::KeyContainsNull => { write!(f, "KVP key must not contain null bytes") @@ -39,6 +101,12 @@ impl fmt::Display for KvpError { Self::MaxUniqueKeysExceeded { max } => { write!(f, "KVP unique key count exceeded maximum ({max})") } + Self::PayloadNotUtf8 => { + write!(f, "diagnostic payload must be valid UTF-8 for encoding 'none'") + } + Self::UnsupportedEncoding { token } => { + write!(f, "diagnostic encoding '{token}' is not supported") + } Self::ValueTooLarge { max, actual } => { write!(f, "KVP value length ({actual}) exceeds maximum ({max})") } @@ -63,3 +131,47 @@ impl From for KvpError { Self::Io(err) } } + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[rstest] + #[case( + KvpError::EmptyEventField { field: "name" }, + "event key field 'name' must not be empty" + )] + #[case( + KvpError::EventFieldContainsDelimiter { field: "name" }, + "event key field 'name' must not contain '|'" + )] + #[case( + KvpError::EventFieldTooLong { field: "name", max: 48, actual: 49 }, + "event key field 'name' length (49) exceeds maximum (48)" + )] + #[case( + KvpError::InvalidUuid { field: "vm_id" }, + "event key field 'vm_id' must be a UUID" + )] + #[case( + KvpError::TooManyChunks { max: 1023 }, + "diagnostic chunk count exceeds maximum (1023)" + )] + #[case( + KvpError::PayloadNotUtf8, + "diagnostic payload must be valid UTF-8 for encoding 'none'" + )] + #[case( + KvpError::UnsupportedEncoding { token: "zstd+b64".into() }, + "diagnostic encoding 'zstd+b64' is not supported" + )] + fn diagnostic_validation_errors_explain_the_failure( + #[case] error: KvpError, + #[case] expected: &str, + ) { + assert_eq!(error.to_string(), expected); + let error: &dyn std::error::Error = &error; + assert!(error.source().is_none()); + } +} diff --git a/libazureinit-kvp/src/lib.rs b/libazureinit-kvp/src/lib.rs index f17b66f7..32c040d5 100644 --- a/libazureinit-kvp/src/lib.rs +++ b/libazureinit-kvp/src/lib.rs @@ -1,22 +1,80 @@ // Copyright (c) Microsoft Corporation. // Licensed under the MIT License. -//! `libazureinit-kvp` provides a unified KVP pool file store for -//! Hyper-V/Azure guests. -//! -//! - [`KvpPoolStore`]: KVP pool file store with -//! [`PoolMode`]-based policy. -//! - [`ProvisioningReport`]: structured provisioning health report that -//! is persisted as the single `PROVISIONING_REPORT` record with -//! [`write_report`]. +//! Read and write Hyper-V KVP metadata on Linux guests. +//! +//! KVP (Key-Value Pair) data exchange lets a VM share small text records with +//! its host without network connectivity. This library operates on the local +//! pool files, normally in `/var/lib/hyperv`. The Hyper-V daemon and kernel +//! handle transport; writing a file does not itself notify the host. +//! +//! # Choose an API +//! - [`DiagnosticWriter`] emits operations and observations, handling +//! timestamps, compression and splitting long payloads into records. +//! - [`DiagnosticReader`] reads native and cloud-init diagnostics, provisioning +//! reports, and unrecognized records as [`Entry`] values. +//! - [`ProvisioningReport`] and [`write_report`] publish one replaceable +//! provisioning result, rather than a history of diagnostic events. +//! - [`KvpPoolStore`] reads and writes ordinary key/value records. Use +//! [`insert`](KvpPoolStore::insert) to update a key, or +//! [`append`](KvpPoolStore::append) to preserve duplicates. +//! +//! Use [`KvpPool::Guest`] for guest-produced telemetry and [`PoolMode::Safe`] +//! for host-readable writes. [`KvpPool::AutoExternal`] contains host-provided +//! data. The selected directory must exist and permit the requested file I/O; +//! [`KvpPoolStore::new_in`] selects an alternate directory. +//! +//! # Emit Telemetry +//! +//! ```no_run +//! use libazureinit_kvp::{ +//! DiagnosticWriter, Encoding, KvpPool, KvpPoolStore, Outcome, PoolMode, +//! }; +//! +//! # fn main() -> Result<(), Box> { +//! let store = KvpPoolStore::new(KvpPool::Guest, PoolMode::Safe)?; +//! let writer = DiagnosticWriter::new( +//! store, +//! "azure-init/0.1.1", +//! "3f2504e0-4f89-41d3-9a0c-0305e82c3301", +//! )?; +//! writer.emit_event( +//! "imds", "metadata retrieved", None, Some(Outcome::Success), None, +//! )?; +//! writer.emit_event( +//! "os:release", std::fs::read("/etc/os-release")?, +//! Some(Encoding::ZlibB64), None, None, +//! )?; +//! # Ok(()) +//! # } +//! ``` +//! +//! See [`DiagnosticWriter`] for recording operations, [`DiagnosticReader`] for +//! consuming telemetry, and [`ProvisioningReport`] for reporting provisioning +//! success or failure. +//! +//! # Format References +//! The [KVP contract] describes pool files and Hyper-V interfaces. The +//! [diagnostics contract] describes record fields and encodings for consumers +//! that read telemetry independently of this library. +//! +//! [KVP contract]: https://github.com/Azure/azure-init/blob/main/doc/kvp.md +//! [diagnostics contract]: https://github.com/Azure/azure-init/blob/main/doc/diagnostics.md mod cli; +mod diagnostics; mod error; mod report; mod store; mod vm_id; pub use cli::run; +pub use diagnostics::{ + DecodeError, Diagnostic, DiagnosticEvent, DiagnosticFinish, DiagnosticKey, + DiagnosticPayload, DiagnosticReader, DiagnosticStart, DiagnosticWriter, + DurationPrecision, Encoding, Entry, Kind, Outcome, RawKeyValue, + TimestampPrecision, DIAGNOSTIC_VERSION_ID, 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..3690c3b7 100644 --- a/libazureinit-kvp/src/report.rs +++ b/libazureinit-kvp/src/report.rs @@ -1,34 +1,24 @@ // Copyright (c) Microsoft Corporation. // Licensed under the MIT License. -//! Structured provisioning report abstraction layered over the raw -//! [`KvpPoolStore`](crate::KvpPoolStore) key/value API. -//! -//! [`ProvisioningReport`] is a strongly-typed representation of a -//! provisioning health report. Instead of building ad-hoc key/value -//! strings at the call site, callers construct a report and persist it -//! with [`write_report`], which serializes it into the single -//! pipe-delimited `PROVISIONING_REPORT` KVP record that the Azure/Hyper-V -//! host parses. +//! Provisioning report creation and storage. -use chrono::Utc; +use std::str::FromStr; -use crate::{KvpError, KvpPoolStore}; +use chrono::{DateTime, Utc}; -/// KVP key under which the encoded provisioning health report is stored. -/// -/// The Azure/Hyper-V host parses this single key; its value is the -/// pipe-delimited `key=value|key=value|...` report produced by -/// [`write_report`]. +use crate::{DecodeError, KvpError, KvpPoolStore}; + +/// Key used by [`write_report`] to store the provisioning result. pub const PROVISIONING_REPORT_KEY: &str = "PROVISIONING_REPORT"; -/// The current time formatted as an RFC 3339 string. fn now_rfc3339() -> String { Utc::now().to_rfc3339() } /// Outcome of a provisioning attempt. -#[derive(Clone, Copy, Debug, PartialEq, Eq)] +#[derive(Clone, Copy, Debug, PartialEq, Eq, serde::Serialize)] +#[serde(rename_all = "lowercase")] enum ReportResult { /// Provisioning completed successfully. Success, @@ -52,27 +42,34 @@ 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. -#[derive(Clone, Copy, Debug, PartialEq, Eq)] +/// Pre-provisioning state included in a provisioning report. +#[derive(Clone, Copy, Debug, PartialEq, Eq, serde::Serialize)] pub enum ReportPpsType { - /// Not pre-provisioned (`None`). + /// Not pre-provisioned. None, - /// Pre-provisioned OS disk (`PreprovisionedOSDisk`). + /// A pre-provisioned OS disk. + #[serde(rename = "PreprovisionedOSDisk")] OsDisk, - /// Running pre-provisioning (`Running`). + /// Pre-provisioning is running. Running, - /// Savable pre-provisioning (`Savable`). + /// Pre-provisioning can be saved. Savable, - /// Unknown pre-provisioning type (`Unknown`). + /// The pre-provisioning state is unknown. Unknown, } impl ReportPpsType { + fn from_wire(value: &str) -> Result { + match value { + "None" => Ok(Self::None), + "PreprovisionedOSDisk" => Ok(Self::OsDisk), + "Running" => Ok(Self::Running), + "Savable" => Ok(Self::Savable), + "Unknown" => Ok(Self::Unknown), + _ => Err(DecodeError::Malformed), + } + } + /// The wire string used in the `pps_type` KVP field. fn as_str(self) -> &'static str { match self { @@ -91,11 +88,16 @@ impl std::fmt::Display for ReportPpsType { } } -/// A strongly-typed provisioning health report. +/// A provisioning result for host telemetry. +/// +/// [`success`](Self::success) and [`failure`](Self::failure) capture the current +/// time. Add optional context with the builder methods, then call +/// [`write_report`] to persist it. Existing stored reports can be parsed with +/// [`str::parse`]; parsing preserves their timestamps. +/// +/// See the [provisioning report contract] for the stored format. /// -/// Construct one with [`ProvisioningReport::success`] or -/// [`ProvisioningReport::failure`], optionally attach extra context with -/// the builder methods, then persist it with [`write_report`]. +/// [provisioning report contract]: https://github.com/Azure/azure-init/blob/main/doc/diagnostics.md#provisioning-reports /// /// # Example /// ```no_run @@ -118,29 +120,31 @@ impl std::fmt::Display for ReportPpsType { /// # Ok(()) /// # } /// ``` -#[derive(Clone, Debug, PartialEq, Eq)] +#[derive(Clone, Debug, PartialEq, Eq, serde::Serialize)] pub struct ProvisioningReport { - /// Provisioning outcome (`result` field). + /// Provisioning outcome. result: ReportResult, - /// Reporting agent identifier (`agent` field). + /// Reporting agent identifier. agent: String, - /// Virtual machine identifier (`vm_id` field). + /// Virtual machine identifier. vm_id: String, - /// Report timestamp (`timestamp` field), set to the current time - /// (RFC 3339) when the report is constructed. + /// RFC 3339 timestamp captured when the report is constructed. timestamp: String, - /// Pre-provisioning type (`pps_type` field). + /// Pre-provisioning state. pps_type: ReportPpsType, - /// Failure reason (`reason` field). Present for error reports. + /// Failure reason, present only for error reports. + #[serde(skip_serializing_if = "Option::is_none")] reason: Option, - /// Documentation URL (`documentation_url` field), if applicable. + /// Help URL, if applicable. + #[serde(skip_serializing_if = "Option::is_none")] documentation_url: Option, - /// Additional ordered key/value context (e.g. supporting data). + /// Additional ordered key/value context. + #[serde(skip_serializing_if = "Vec::is_empty")] extra: Vec<(String, String)>, } impl ProvisioningReport { - /// Create a successful provisioning report. + /// Creates a successful provisioning report. pub fn success( agent: impl Into, vm_id: impl Into, @@ -158,7 +162,7 @@ impl ProvisioningReport { } } - /// Create a failed provisioning report with a failure reason. + /// Creates a failed provisioning report with a reason. pub fn failure( agent: impl Into, vm_id: impl Into, @@ -177,14 +181,16 @@ impl ProvisioningReport { } } - /// Attach a documentation URL. + /// Sets a help URL, included in the stored value for failure reports only. pub fn with_documentation_url(mut self, url: impl Into) -> Self { self.documentation_url = Some(url.into()); self } - /// Append an additional key/value pair. Extras are emitted in the - /// order they were added. + /// Adds context to the report. + /// + /// Entries retain insertion order and duplicate keys. Do not use standard + /// report field names such as `result` or `timestamp`. pub fn with_extra( mut self, key: impl Into, @@ -195,15 +201,130 @@ impl ProvisioningReport { } } +impl FromStr for ProvisioningReport { + type Err = DecodeError; + + /// Parses a stored report without changing its timestamp. + fn from_str(value: &str) -> Result { + let value = value + .strip_suffix("\r\n") + .or_else(|| value.strip_suffix('\n')) + .unwrap_or(value); + validate_report_quoting(value)?; + let mut reader = csv::ReaderBuilder::new() + .delimiter(b'|') + .has_headers(false) + .from_reader(value.as_bytes()); + let record = reader + .records() + .next() + .ok_or(DecodeError::Malformed)? + .map_err(|_| DecodeError::Malformed)?; + let mut fields = record + .iter() + .map(|field| { + let (key, value) = + field.split_once('=').ok_or(DecodeError::Malformed)?; + if key.is_empty() { + return Err(DecodeError::Malformed); + } + Ok((key.to_owned(), value.to_owned())) + }) + .collect::, DecodeError>>()?; + + let result = match take_field(&mut fields, "result")? + .ok_or(DecodeError::Malformed)? + .as_str() + { + "success" => ReportResult::Success, + "error" => ReportResult::Error, + _ => return Err(DecodeError::Malformed), + }; + let agent = + take_field(&mut fields, "agent")?.ok_or(DecodeError::Malformed)?; + let vm_id = + take_field(&mut fields, "vm_id")?.ok_or(DecodeError::Malformed)?; + let timestamp = take_field(&mut fields, "timestamp")? + .ok_or(DecodeError::Malformed)?; + DateTime::parse_from_rfc3339(×tamp) + .map_err(|_| DecodeError::Malformed)?; + let pps_type = ReportPpsType::from_wire( + &take_field(&mut fields, "pps_type")? + .ok_or(DecodeError::Malformed)?, + )?; + let (reason, documentation_url) = match result { + ReportResult::Success => (None, None), + ReportResult::Error => ( + Some( + take_field(&mut fields, "reason")? + .ok_or(DecodeError::Malformed)?, + ), + take_field(&mut fields, "documentation_url")?, + ), + }; + Ok(Self { + result, + agent, + vm_id, + timestamp, + pps_type, + reason, + documentation_url, + extra: fields, + }) + } +} + +fn take_field( + fields: &mut Vec<(String, String)>, + name: &str, +) -> Result, DecodeError> { + let mut matches = fields + .iter() + .enumerate() + .filter(|(_, (key, _))| key == name); + let position = matches.next().map(|(index, _)| index); + if matches.next().is_some() { + return Err(DecodeError::Malformed); + } + Ok(position.map(|index| fields.remove(index).1)) +} + +fn validate_report_quoting(value: &str) -> Result<(), DecodeError> { + enum State { + Start, + Unquoted, + Quoted, + Closed, + } + use State::*; + + // The CSV reader tolerates broken quoting, but reports must be unambiguous. + let mut state = Start; + for byte in value.bytes() { + if byte == 0 { + return Err(DecodeError::Malformed); + } + state = match (state, byte) { + (Start | Closed, b'"') => Quoted, + (Quoted, b'"') => Closed, + (Quoted, _) => Quoted, + (Start | Unquoted | Closed, b'|') => Start, + (Closed, _) | (_, b'"' | b'\r' | b'\n') => { + return Err(DecodeError::Malformed); + } + _ => Unquoted, + }; + } + if matches!(state, Quoted) { + return Err(DecodeError::Malformed); + } + Ok(()) +} + impl ProvisioningReport { - /// Encode the report as a single pipe-delimited `key=value` string. - /// - /// - Success: `result`, `agent`, `pps_type`, `vm_id`, `timestamp`, - /// then any extras in insertion order. - /// - Failure: `result`, `reason`, `agent`, extras in insertion - /// order, `pps_type`, `vm_id`, `timestamp`, then - /// `documentation_url` (if any). - fn encode(&self) -> String { + /// Encodes the report for storage. + pub(crate) fn encode(&self) -> String { let mut data = Vec::with_capacity(7 + self.extra.len()); data.push(format!("result={}", self.result)); @@ -251,12 +372,12 @@ impl ProvisioningReport { } } -/// Persist a report to the KVP store under [`PROVISIONING_REPORT_KEY`]. +/// Stores a provisioning result, replacing any existing report in the pool. +/// +/// A report must fit in one value under the store's configured size policy. /// -/// The report is encoded into a single pipe-delimited value and written -/// with [`KvpPoolStore::insert`] (upsert / last-write-wins), so it -/// overrides any existing `PROVISIONING_REPORT` record rather than -/// accumulating duplicates. +/// # Errors +/// Returns validation and I/O errors from the store. pub fn write_report( store: &KvpPoolStore, report: &ProvisioningReport, @@ -286,6 +407,15 @@ mod tests { report } + fn success_wire() -> String { + with_ts(ProvisioningReport::success( + AGENT, + VM_ID, + ReportPpsType::None, + )) + .encode() + } + #[rstest] #[case::success( with_ts(ProvisioningReport::success(AGENT, VM_ID, ReportPpsType::None)), @@ -328,11 +458,12 @@ mod tests { )), "result=error|reason=boom|agent=Azure-Init/0.0.0|pps_type=None|vm_id=00000000-0000-0000-0000-000000000abc|timestamp=2026-06-17T00:00:00+00:00", )] - fn encode_emits_expected_pipe_string( + fn report_wire_format_round_trips( #[case] report: ProvisioningReport, #[case] expected: &str, ) { assert_eq!(report.encode(), expected); + assert_eq!(expected.parse::().unwrap(), report); } /// Pins each [`ReportPpsType`] variant to its exact wire string. @@ -342,11 +473,65 @@ mod tests { #[case(ReportPpsType::Running, "Running")] #[case(ReportPpsType::Savable, "Savable")] #[case(ReportPpsType::Unknown, "Unknown")] - fn pps_type_display_matches_wire_string( + fn pps_type_wire_tokens_match_the_model( #[case] pps_type: ReportPpsType, #[case] expected: &str, ) { assert_eq!(pps_type.to_string(), expected); + assert_eq!(serde_json::to_value(pps_type).unwrap(), expected); + assert_eq!(ReportPpsType::from_wire(expected).unwrap(), pps_type); + } + + #[test] + fn success_serializes_without_absent_fields() { + let report = with_ts(ProvisioningReport::success( + AGENT, + VM_ID, + ReportPpsType::None, + )); + assert_eq!( + serde_json::to_value(report).unwrap(), + serde_json::json!({ + "result": "success", + "agent": AGENT, + "vm_id": VM_ID, + "timestamp": TS, + "pps_type": "None", + }) + ); + } + + #[test] + fn failure_serialization_preserves_ordered_extras() { + let report = with_ts( + ProvisioningReport::failure( + AGENT, + VM_ID, + "boom", + ReportPpsType::OsDisk, + ) + .with_extra("detail", "first") + .with_extra("detail", "second") + .with_extra("result", "extra context") + .with_documentation_url("https://aka.ms/linuxprovisioningerror"), + ); + assert_eq!( + serde_json::to_value(report).unwrap(), + serde_json::json!({ + "result": "error", + "agent": AGENT, + "vm_id": VM_ID, + "timestamp": TS, + "pps_type": "PreprovisionedOSDisk", + "reason": "boom", + "documentation_url": "https://aka.ms/linuxprovisioningerror", + "extra": [ + ["detail", "first"], + ["detail", "second"], + ["result", "extra context"], + ], + }) + ); } /// The success layout lists the standard fields first, then any @@ -454,4 +639,180 @@ mod tests { let result = write_report(&store, &report); assert!(result.is_err()); } + + #[test] + fn parses_reordered_fields_without_normalizing_identity_or_timestamp() { + let timestamp = "2026-06-17T02:00:00.123456+02:00"; + let value = format!( + "timestamp={timestamp}|vm_id=vm-abc|pps_type=None|agent={AGENT}|result=success" + ); + let mut expected = + ProvisioningReport::success(AGENT, "vm-abc", ReportPpsType::None); + expected.timestamp = timestamp.into(); + assert_eq!(value.parse::().unwrap(), expected); + } + + #[test] + fn failure_round_trips_quoted_fields_and_documentation_url() { + let report = with_ts( + ProvisioningReport::failure( + AGENT, + VM_ID, + "failed | \"quoted\"\r\nnext line", + ReportPpsType::Running, + ) + .with_extra("context", "key=value|details\nmore") + .with_documentation_url("https://example.invalid/?key=a=b"), + ); + assert_eq!( + report.encode().parse::().unwrap(), + report + ); + } + + #[test] + fn supporting_data_preserves_order_duplicates_and_empty_values() { + let value = + format!("{}|detail=first|detail=second|empty=", success_wire()); + let report = value.parse::().unwrap(); + assert_eq!( + report.extra, + vec![ + ("detail".into(), "first".into()), + ("detail".into(), "second".into()), + ("empty".into(), String::new()), + ] + ); + } + + #[test] + fn success_keeps_failure_only_fields_as_supporting_data() { + let report = with_ts( + ProvisioningReport::success(AGENT, VM_ID, ReportPpsType::None) + .with_extra("reason", "additional context") + .with_extra("documentation_url", "https://example.invalid/"), + ); + assert_eq!( + report.encode().parse::().unwrap(), + report + ); + } + + #[rstest] + #[case::result("result")] + #[case::agent("agent")] + #[case::vm_id("vm_id")] + #[case::timestamp("timestamp")] + #[case::pps_type("pps_type")] + fn report_requires_each_standard_field(#[case] missing: &str) { + let value = success_wire() + .split('|') + .filter(|field| !field.starts_with(&format!("{missing}="))) + .collect::>() + .join("|"); + assert_eq!( + value.parse::(), + Err(DecodeError::Malformed) + ); + } + + #[test] + fn failure_requires_a_reason() { + let value = success_wire().replace("result=success", "result=error"); + assert_eq!( + value.parse::(), + Err(DecodeError::Malformed) + ); + } + + #[rstest] + #[case::success(with_ts(ProvisioningReport::success( + "", + "", + ReportPpsType::None + )))] + #[case::failure(with_ts(ProvisioningReport::failure( + AGENT, + VM_ID, + "", + ReportPpsType::None + )))] + fn empty_values_supported_by_the_writer_remain_readable( + #[case] report: ProvisioningReport, + ) { + assert_eq!( + report.encode().parse::().unwrap(), + report + ); + } + + #[rstest] + #[case::result("result=success", "result=fail")] + #[case::pps_type("pps_type=None", "pps_type=FutureType")] + #[case::timestamp(TS, "not-a-timestamp")] + fn invalid_standard_values_are_malformed( + #[case] from: &str, + #[case] to: &str, + ) { + let value = success_wire().replace(from, to); + assert_eq!( + value.parse::(), + Err(DecodeError::Malformed) + ); + } + + #[rstest] + #[case::unclosed_quote("\"extra=value")] + #[case::characters_after_quote("\"extra=value\"suffix")] + #[case::unquoted_quote("extra=va\"lue")] + #[case::missing_equals("extra")] + #[case::empty_key("=value")] + #[case::null("extra=va\0lue")] + fn malformed_supporting_fields_are_not_silently_repaired( + #[case] extra: &str, + ) { + let value = format!("{}|{extra}", success_wire()); + assert_eq!( + value.parse::(), + Err(DecodeError::Malformed) + ); + } + + #[test] + fn multiple_csv_records_are_rejected() { + let value = format!("{}\n{}", success_wire(), success_wire()); + assert_eq!( + value.parse::(), + Err(DecodeError::Malformed) + ); + } + + #[test] + fn optional_csv_quotes_and_record_terminators_are_accepted() { + let quoted = success_wire() + .split('|') + .map(|field| format!("\"{field}\"")) + .collect::>() + .join("|"); + let expected = success_wire().parse::().unwrap(); + for ending in ["", "\n", "\r\n"] { + assert_eq!( + format!("{quoted}{ending}") + .parse::() + .unwrap(), + expected + ); + } + } + + #[rstest] + #[case::conflicting_result("result=error")] + #[case::repeated_agent("agent=Azure-Init/0.0.0")] + fn duplicate_standard_fields_are_ambiguous(#[case] duplicate: &str) { + let value = format!("{}|{duplicate}", success_wire()); + assert_eq!( + value.parse::(), + Err(DecodeError::Malformed) + ); + } } diff --git a/libazureinit-kvp/src/store.rs b/libazureinit-kvp/src/store.rs index a067c980..de06c40e 100644 --- a/libazureinit-kvp/src/store.rs +++ b/libazureinit-kvp/src/store.rs @@ -1,26 +1,7 @@ // Copyright (c) Microsoft Corporation. // Licensed under the MIT License. -//! Unified KVP pool file backend for Hyper-V and Azure guests. -//! -//! All sizes are UTF-8 byte counts — the on-disk pool stores keys and -//! values as zero-padded UTF-8. The Windows-side Hyper-V wire spec is -//! UTF-16 char-based (256 / 1024), but the Linux `hv_kvp_daemon` -//! exchanges UTF-8 bytes with the kernel. -//! -//! Fixed-width record format (matches Linux kernel -//! `HV_KVP_EXCHANGE_MAX_*`): -//! - key field: 512 bytes -//! - value field: 2048 bytes -//! - record size: 2560 bytes -//! -//! [`PoolMode`] selects which size limits are enforced on writes: -//! - [`Safe`](PoolMode::Safe): key <= 254 bytes, -//! value <= 1022 bytes (recommended for Linux kernel compatibility) -//! - [`Unsafe`](PoolMode::Unsafe): key <= 512 bytes, value <= 2048 bytes -//! -//! ## Reference -//! - [Hyper-V Data Exchange Service (KVP)](https://learn.microsoft.com/en-us/virtualization/hyper-v-on-windows/reference/integration-services#hyper-v-data-exchange-service-kvp) +//! KVP pool storage and file locking. use std::collections::{HashMap, HashSet}; use std::fs::{File, OpenOptions}; @@ -40,12 +21,15 @@ const WIRE_MAX_VALUE_BYTES: usize = 2048; const SAFE_MAX_KEY_BYTES: usize = 254; const SAFE_MAX_VALUE_BYTES: usize = 1022; -/// Maximum number of unique keys allowed in the pool. +/// Unique-key limit for insert and load operations. const MAX_UNIQUE_KEYS: usize = 1024; const RECORD_SIZE: usize = WIRE_MAX_KEY_BYTES + WIRE_MAX_VALUE_BYTES; -/// Hyper-V KVP pool indices. +/// Selects the Hyper-V pool file to access. +/// +/// Use [`Guest`](Self::Guest) for guest-produced diagnostics and reports. +/// Other pools are generally populated by the host or integration daemon. #[derive(Clone, Copy, Debug, PartialEq, Eq)] #[repr(u8)] pub enum KvpPool { @@ -53,11 +37,11 @@ pub enum KvpPool { External = 0, /// Guest-to-host data; cloud-init / azure-init write here (`.kvp_pool_1`). Guest = 1, - /// Guest intrinsics generated by the daemon (`.kvp_pool_2`). + /// Guest details generated by the daemon, not read from `.kvp_pool_2`. Auto = 2, /// Host-originated data describing the host (`.kvp_pool_3`). AutoExternal = 3, - /// Undocumented; no pool file exists (`.kvp_pool_4`). + /// Internal pool (`.kvp_pool_4`); not intended for application telemetry. AutoInternal = 4, } @@ -84,22 +68,17 @@ impl KvpPool { } } -/// Policy mode controlling key/value size limits for writes. +/// Size policy for keys and values written through [`KvpPoolStore`]. /// -/// All limits are **UTF-8 byte counts** measured against the on-disk -/// pool format. The Windows-side Hyper-V wire spec expresses limits -/// in UTF-16 code units (256 / 1024), which equal 512 / 2048 bytes -/// when stored as UTF-8 by `hv_kvp_daemon` on Linux. +/// Limits count UTF-8 bytes, not characters or UTF-16 code units. This policy +/// does not restrict reading records written with a different mode. #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub enum PoolMode { - /// Conservative limits for Linux kernel compatibility: - /// key <= 254 UTF-8 bytes, value <= 1022 UTF-8 bytes — 2 bytes - /// under the kernel `HV_KVP_EXCHANGE_MAX_*` maximums to leave room - /// for a NUL terminator on either side of the boundary. + /// Host-compatible limits: keys up to 254 bytes and values up to 1,022 bytes. + /// Recommended for data exchanged with the host. Safe, - /// Full Hyper-V wire-format limits expressed in UTF-8 bytes: - /// key <= 512 bytes (== 256 UTF-16 code units), - /// value <= 2048 bytes (== 1024 UTF-16 code units). + /// Full pool-field widths: keys up to 512 bytes and values up to 2,048 bytes. + /// Such records may be truncated during transport to the host. Unsafe, } @@ -119,7 +98,23 @@ impl PoolMode { } } -/// Unified KVP pool file store. +/// Reads and writes UTF-8 key/value records in one Linux Hyper-V pool file. +/// +/// Constructors select a path and size policy without performing I/O. Each +/// operation opens and locks the file as needed. The parent directory must +/// exist before writing. +/// +/// [`insert`](Self::insert) updates a key, while [`append`](Self::append) +/// preserves duplicates. Use [`dump`](Self::dump) for every physical record +/// in pool order, or [`entries`](Self::entries) for the last value of each key. +/// [`load`](Self::load) replaces the entire pool. +/// +/// Records persist across boots until explicitly cleared. Writes reject empty +/// keys and NUL bytes. An I/O failure during a write may leave partial changes. +/// +/// See the [KVP contract] for record layout and transport details. +/// +/// [KVP contract]: https://github.com/Azure/azure-init/blob/main/doc/kvp.md #[derive(Clone, Debug)] pub struct KvpPoolStore { pool: KvpPool, @@ -129,12 +124,10 @@ pub struct KvpPoolStore { } impl KvpPoolStore { - /// Append one key-value pair to the tail of the pool without - /// checking for an existing key. + /// Appends a key/value pair, preserving existing records and duplicates. /// - /// This preserves any existing records, including duplicate keys, - /// and does not enforce [`MAX_UNIQUE_KEYS`]. Use - /// [`insert`](Self::insert) when callers need upsert semantics. + /// Record count is not capped. Use [`insert`](Self::insert) to replace a + /// key's existing values. pub fn append(&self, key: &str, value: &str) -> Result<(), KvpError> { validate_key(key, self.mode.max_key_size())?; validate_value(value, self.mode.max_value_size())?; @@ -148,18 +141,12 @@ impl KvpPoolStore { Ok(()) } - /// Append records to the tail of the pool under one exclusive - /// lock, preserving order and preventing cooperating writers from - /// interleaving with the batch. + /// Appends records in iteration order under one exclusive lock. /// - /// 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 - /// 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 - /// error occurs after writing starts, the file may contain a - /// partial batch. + /// Existing records and duplicates are kept, and record count is not + /// capped. Empty input is a no-op. Inputs are validated before writing. + /// Cooperating writers cannot interleave with the batch; an I/O error may + /// leave only part of it written. pub fn append_multiple(&self, records: I) -> Result<(), KvpError> where I: IntoIterator, @@ -185,7 +172,7 @@ impl KvpPoolStore { self.write_records_locked(&records, false) } - /// Remove all entries from the store. + /// Removes all records. A missing pool file is left absent. pub fn clear(&self) -> Result<(), KvpError> { { let mut handle = match self.open_for_read_write() { @@ -200,7 +187,9 @@ impl KvpPoolStore { Ok(()) } - /// Clear the pool file if it contains data from a previous boot. + /// Clears the pool if its modification time is at or before system boot. + /// + /// The check and removal hold an exclusive lock. A missing file is a no-op. pub fn clear_if_stale(&self) -> Result<(), KvpError> { { let mut handle = match self.open_for_read_write() { @@ -218,12 +207,12 @@ impl KvpPoolStore { Ok(()) } - /// Remove all records matching the key. Returns `true` if at least - /// one record was present. + /// Removes all records matching `key`, returning whether any were found. /// /// Rejects empty keys and keys containing null bytes; key size is /// not capped, so a [`Safe`](PoolMode::Safe)-mode store can delete /// keys written in [`Unsafe`](PoolMode::Unsafe) mode. + /// Removal may change the order of remaining records. pub fn delete(&self, key: &str) -> Result { validate_key(key, usize::MAX)?; @@ -252,16 +241,13 @@ impl KvpPoolStore { Ok(found) } - /// Delete all records whose key appears in `keys`, under one - /// exclusive lock. + /// Removes all records whose key appears in `keys` under one exclusive lock. /// - /// Returns the number of records removed, counting duplicate - /// on-disk keys separately. Like [`delete`](Self::delete), key - /// size is not capped. Empty input is a no-op; a missing file - /// returns `Ok(0)`. + /// Returns the number of records removed, counting duplicates separately. + /// Like [`delete`](Self::delete), key size is not capped. Empty input is a + /// no-op; a missing file returns `Ok(0)`. /// - /// Removal swaps deleted records with the tail, so remaining - /// record order is not preserved. + /// Remaining records may be reordered. pub fn delete_multiple(&self, keys: I) -> Result where I: IntoIterator, @@ -303,8 +289,10 @@ impl KvpPoolStore { Ok(removed) } - /// Return all key-value records in on-disk order, including - /// duplicates from [`append`](Self::append) calls. + /// Returns every key/value record in pool order, including duplicates. + /// + /// A missing file returns an empty list. Invalid record layout or UTF-8, + /// or a file-access error, fails the call without returning partial results. pub fn dump(&self) -> Result, KvpError> { let records = { let mut iter = match self.iter() { @@ -324,7 +312,10 @@ impl KvpPoolStore { Ok(records) } - /// Return all key-value pairs (deduplicated, last-write-wins). + /// Returns one value per key, using the last record when keys repeat. + /// + /// A missing file returns an empty map. Use [`dump`](Self::dump) to retain + /// duplicate records and their order. pub fn entries(&self) -> Result, KvpError> { let map = { let mut iter = match self.iter() { @@ -345,7 +336,10 @@ impl KvpPoolStore { Ok(map) } - /// Insert a new key-value pair or update an existing key's value. + /// Inserts a key/value pair or replaces an existing key's value. + /// + /// Existing duplicates of the key are collapsed into one record. Adding + /// a new key fails if the pool already contains 1,024 distinct keys. pub fn insert(&self, key: &str, value: &str) -> Result<(), KvpError> { validate_key(key, self.mode.max_key_size())?; validate_value(value, self.mode.max_value_size())?; @@ -387,12 +381,13 @@ impl KvpPoolStore { Ok(()) } - /// Return whether the store is empty. + /// Returns whether the store is empty. pub fn is_empty(&self) -> Result { Ok(self.len()? == 0) } - /// Whether the store's data is stale (e.g. predates current boot). + /// Returns whether the pool's modification time is at or before system boot. + /// A missing file is not stale. pub fn is_stale(&self) -> Result { let metadata = match self.ops.path_metadata(&self.path) { Ok(m) => m, @@ -425,12 +420,7 @@ impl KvpPoolStore { KvpPoolIter::new(handle, false) } - /// Return the number of records in the store. - /// - /// This counts on-disk records, not unique keys. If - /// [`append`](Self::append) was used to write duplicate keys, this - /// may exceed the number of unique keys returned by - /// [`entries`](Self::entries). + /// Returns the number of physical records, including duplicate keys. pub fn len(&self) -> Result { match self.ops.path_metadata(&self.path) { Ok(m) => record_count_from_len(m.len as usize).map_err(Into::into), @@ -454,13 +444,15 @@ impl KvpPoolStore { self.mode } - /// Open using the default Hyper-V directory - /// ([`/var/lib/hyperv`](KvpPool::default_dir)). + /// Selects a pool in the default `/var/lib/hyperv` directory without I/O. pub fn new(pool: KvpPool, mode: PoolMode) -> Result { Self::new_in(pool, KvpPool::default_dir(), mode) } - /// Open using a custom directory (file name derived from `pool`). + /// Selects a pool in `dir` without opening or creating it. + /// + /// The file name is derived from [`KvpPool::file_name`]. The directory must + /// exist before performing a write that creates the file. pub fn new_in( pool: KvpPool, dir: impl AsRef, @@ -522,8 +514,7 @@ impl KvpPoolStore { iter.flush()?; Ok(()) } - - /// Return a reference to the pool file path. + /// Returns the configured pool-file path. pub fn path(&self) -> &Path { &self.path } @@ -533,23 +524,16 @@ impl KvpPoolStore { self.pool } - /// Replace the entire pool with the given pairs, in iteration - /// order, under one exclusive lock. + /// Replaces the entire pool in iteration order under one exclusive lock. /// - /// 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; - /// empty input clears the pool. Use - /// [`append_multiple`](Self::append_multiple) when callers need to - /// extend the pool instead. + /// Duplicates are preserved, with at most 1,024 distinct keys. Empty input + /// clears the pool. Use [`append_multiple`](Self::append_multiple) to keep + /// existing records instead. /// - /// Validation happens before locking, so a rejected batch never - /// blocks other writers. The file is truncated and rewritten - /// under one exclusive lock; an I/O error mid-write may leave - /// the file partially written. Every successful call updates - /// mtime (relevant to [`is_stale`](Self::is_stale)). Rejects - /// malformed pool files (size not a multiple of the record - /// size); call [`clear`](Self::clear) first to recover. + /// Invalid input leaves the pool unchanged; a storage error after truncation + /// may leave a partial replacement. A successful call updates modification + /// time. An existing file with invalid record framing must be explicitly + /// cleared with [`clear`](Self::clear) before it can be replaced. pub fn load(&self, records: I) -> Result<(), KvpError> where I: IntoIterator, @@ -578,10 +562,7 @@ impl KvpPoolStore { self.write_records_locked(&records, true) } - /// Read the value for a key. Returns `Ok(None)` when absent. - /// - /// If multiple records share the same key (e.g. via - /// [`append`](Self::append)), the last (most recent) match wins. + /// Returns the last value for `key`, or `None` when absent. /// /// Rejects empty keys and keys containing null bytes; key size is /// not capped, so a [`Safe`](PoolMode::Safe)-mode store can read @@ -653,13 +634,19 @@ fn decode_record(data: &[u8]) -> io::Result<(String, String)> { let (key_bytes, value_bytes) = data.split_at(WIRE_MAX_KEY_BYTES); - let key = std::str::from_utf8(key_bytes) + let key_end = key_bytes + .iter() + .position(|byte| *byte == 0) + .unwrap_or(key_bytes.len()); + let key = std::str::from_utf8(&key_bytes[..key_end]) .map_err(|e| io::Error::new(ErrorKind::InvalidData, e))? - .trim_end_matches('\0') .to_string(); - let value = std::str::from_utf8(value_bytes) + let value_end = value_bytes + .iter() + .position(|byte| *byte == 0) + .unwrap_or(value_bytes.len()); + let value = std::str::from_utf8(&value_bytes[..value_end]) .map_err(|e| io::Error::new(ErrorKind::InvalidData, e))? - .trim_end_matches('\0') .to_string(); Ok((key, value)) @@ -1234,6 +1221,18 @@ mod tests { KvpErrKind::ValueContainsNull )] #[case::load_empty_key(WriteOp::Load, "", "v", KvpErrKind::EmptyKey)] + #[case::append_multiple_empty_key( + WriteOp::AppendMultiple, + "", + "v", + KvpErrKind::EmptyKey + )] + #[case::append_multiple_null_key( + WriteOp::AppendMultiple, + "bad\0key", + "v", + KvpErrKind::KeyContainsNull + )] fn test_write_rejects_invalid_input( #[case] op: WriteOp, #[case] key: &str, @@ -1246,9 +1245,13 @@ mod tests { WriteOp::Insert => store.insert(key, value), WriteOp::Append => store.append(key, value), WriteOp::Load => store.load(pairs([(key, value)])), + WriteOp::AppendMultiple => { + store.append_multiple(pairs([(key, value)])) + } } .unwrap_err(); assert!(expected.matches(&err), "got: {err:?}"); + assert!(!store.path().exists(), "no file on invalid input"); } #[derive(Clone, Copy)] @@ -1268,6 +1271,7 @@ mod tests { Insert, Append, Load, + AppendMultiple, } #[derive(Clone, Copy)] @@ -1359,7 +1363,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); } @@ -1585,6 +1588,16 @@ mod tests { "bad\0key", KvpErrKind::KeyContainsNull )] + #[case::delete_multiple_empty( + ReadOp::DeleteMultiple, + "", + KvpErrKind::EmptyKey + )] + #[case::delete_multiple_null( + ReadOp::DeleteMultiple, + "bad\0key", + KvpErrKind::KeyContainsNull + )] fn test_bad_key_is_rejected( #[case] op: ReadOp, #[case] bad_key: &str, @@ -1597,6 +1610,9 @@ mod tests { let err = match op { ReadOp::Read => store.read(bad_key).unwrap_err(), ReadOp::Delete => store.delete(bad_key).unwrap_err(), + ReadOp::DeleteMultiple => { + store.delete_multiple(vec![bad_key]).unwrap_err() + } }; assert!(expected.matches(&err), "got {err:?}"); assert_eq!(store.read("k1").unwrap(), Some("v1".to_string())); @@ -1607,6 +1623,7 @@ mod tests { enum ReadOp { Read, Delete, + DeleteMultiple, } #[test] @@ -1660,7 +1677,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 +1770,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 +1794,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 +1802,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 +1826,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 +1861,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 +1878,6 @@ mod tests { .append_multiple(Vec::<(String, String)>::new()) .unwrap(); - // No file created when the input is empty. assert!(!store.path().exists()); } @@ -1889,46 +1897,16 @@ 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")])); } - #[test] - fn test_append_multiple_rejects_empty_key() { - let dir = TempDir::new().unwrap(); - let store = safe_store(dir.path()); - - let err = store - .append_multiple(pairs([("ok", "v"), ("", "bad")])) - .unwrap_err(); - assert!(matches!(err, KvpError::EmptyKey), "got {err:?}"); - assert!(!store.path().exists()); - } - - #[test] - fn test_append_multiple_rejects_null_in_key() { - let dir = TempDir::new().unwrap(); - let store = safe_store(dir.path()); - - let err = store - .append_multiple(pairs([("ok\0bad", "v")])) - .unwrap_err(); - assert!(matches!(err, KvpError::KeyContainsNull), "got {err:?}"); - } - #[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 +1936,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 +1984,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,43 +2004,13 @@ 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")])); } - #[test] - fn test_delete_multiple_rejects_empty_key() { - let dir = TempDir::new().unwrap(); - let store = safe_store(dir.path()); - store.load(pairs([("a", "1")])).unwrap(); - - let err = store - .delete_multiple(vec!["a".to_string(), "".to_string()]) - .unwrap_err(); - assert!(matches!(err, KvpError::EmptyKey), "got {err:?}"); - - // Validation runs before any record is removed. - assert_eq!(store.dump().unwrap(), pairs([("a", "1")])); - } - - #[test] - fn test_delete_multiple_rejects_null_in_key() { - let dir = TempDir::new().unwrap(); - let store = safe_store(dir.path()); - store.load(pairs([("a", "1")])).unwrap(); - - let err = store.delete_multiple(vec!["bad\0key"]).unwrap_err(); - assert!(matches!(err, KvpError::KeyContainsNull), "got {err:?}"); - assert_eq!(store.dump().unwrap(), pairs([("a", "1")])); - } - #[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); @@ -2160,6 +2106,25 @@ mod tests { assert_eq!(v, "val"); } + #[rstest] + #[case::key(Field::Key)] + #[case::value(Field::Value)] + fn test_decode_ignores_bytes_after_null_terminator(#[case] field: Field) { + let mut record = encode_record("key", "value"); + let tail = match field { + Field::Key => &mut record[4..WIRE_MAX_KEY_BYTES], + Field::Value => { + &mut record[WIRE_MAX_KEY_BYTES + 6 + ..WIRE_MAX_KEY_BYTES + WIRE_MAX_VALUE_BYTES] + } + }; + tail[..4].copy_from_slice(&[b'x', b'y', 0xFF, 0xFE]); + + let (key, value) = decode_record(&record).unwrap(); + assert_eq!(key, "key"); + assert_eq!(value, "value"); + } + /// Malformed buffers fed to `decode_record` produce the expected /// I/O error kind. #[rstest] @@ -3063,21 +3028,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 +3303,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 +3401,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); diff --git a/libazureinit-kvp/src/vm_id.rs b/libazureinit-kvp/src/vm_id.rs index a8c1edbd..88aafac1 100644 --- a/libazureinit-kvp/src/vm_id.rs +++ b/libazureinit-kvp/src/vm_id.rs @@ -1,22 +1,18 @@ // Copyright (c) Microsoft Corporation. // Licensed under the MIT License. -//! Self-contained VM ID lookup used to auto-populate provisioning reports. +//! Current VM ID lookup for command-line defaults. //! -//! The VM ID is read from `/sys/class/dmi/id/product_uuid` and, on Gen1 VMs, -//! the first three UUID fields are byte-swapped from big-endian to native endianness. +//! Reads `/sys/class/dmi/id/product_uuid` and adjusts Gen1 UUID byte order. use std::fs; use std::path::Path; use uuid::Uuid; -/// Retrieves the current VM ID by reading `/sys/class/dmi/id/product_uuid` -/// and byte-swapping the result if the VM is Gen1. +/// Returns the current VM ID, adjusting UUID byte order on Gen1 VMs. /// -/// # Returns -/// - `Some(String)` containing the VM ID if retrieval is successful. -/// - `None` if the file is missing, empty, or cannot be read. +/// Returns `None` if the DMI file is missing, unreadable or empty. pub fn get_vm_id() -> Option { private_get_vm_id(None, None, None) } @@ -58,8 +54,7 @@ fn private_get_vm_id( } } -/// Determines whether the VM is Gen1 (i.e. not UEFI/Gen2) based on EFI -/// detection. Returns `true` when neither EFI path exists. +/// Returns `true` when neither EFI path exists, identifying a Gen1 VM. fn is_vm_gen1( sysfs_efi_path: Option<&str>, dev_efi_path: Option<&str>, @@ -67,7 +62,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 +171,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 +183,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..a5642641 100644 --- a/libazureinit-kvp/tests/cli.rs +++ b/libazureinit-kvp/tests/cli.rs @@ -4,9 +4,18 @@ use std::fs; use std::io::Write; use std::process::{Command, Output}; - +use std::time::Duration; + +use libazureinit_kvp::{ + DiagnosticWriter, Encoding, KvpPool, KvpPoolStore, Outcome, PoolMode, + PROVISIONING_REPORT_KEY, +}; +use rstest::rstest; +use serde_json::{json, Value}; use tempfile::TempDir; +const VM_ID: &str = "0e5e179d-5341-478b-8456-fbb90621bdf8"; + fn kvp(args: &[&str]) -> Output { Command::new(env!("CARGO_BIN_EXE_libazureinit-kvp")) .args(args) @@ -48,6 +57,14 @@ fn assert_success(output: Output) -> String { String::from_utf8(output.stdout).unwrap() } +fn assert_json(output: Output) -> Value { + serde_json::from_str(&assert_success(output)).unwrap() +} + +fn store_at(dir: &TempDir) -> KvpPoolStore { + KvpPoolStore::new_in(KvpPool::Guest, dir.path(), PoolMode::Safe).unwrap() +} + #[test] fn help_lists_commands() { let stdout = assert_success(kvp(&["--help"])); @@ -114,19 +131,19 @@ fn write_append_read_dump_entries_delete_and_clear() { assert_eq!(assert_success(kvp(&with_dir(&dir, &["read", "a"]))), "2\n"); assert_eq!( - assert_success(kvp(&with_dir(&dir, &["dump"]))), - "a=1\na=2\n" + assert_json(kvp(&with_dir(&dir, &["dump"]))), + json!([{"key": "a", "value": "1"}, {"key": "a", "value": "2"}]) ); assert_eq!(assert_success(kvp(&with_dir(&dir, &["entries"]))), "a=2\n"); assert_eq!( assert_success(kvp(&with_dir(&dir, &["delete", "a"]))), "true\n" ); - assert_eq!(assert_success(kvp(&with_dir(&dir, &["dump"]))), ""); + assert_eq!(assert_json(kvp(&with_dir(&dir, &["dump"]))), json!([])); assert_success(kvp(&with_dir(&dir, &["write", "b", "3"]))); assert_success(kvp(&with_dir(&dir, &["clear"]))); - assert_eq!(assert_success(kvp(&with_dir(&dir, &["dump"]))), ""); + assert_eq!(assert_json(kvp(&with_dir(&dir, &["dump"]))), json!([])); } #[test] @@ -140,7 +157,7 @@ fn load_replaces_pool_from_file() { &["load", "--file", input.to_str().unwrap()], ))); assert_eq!( - assert_success(kvp(&with_dir(&dir, &["dump"]))), + assert_success(kvp(&with_dir(&dir, &["dump", "--text"]))), "a=1\nb=2\n" ); } @@ -166,7 +183,7 @@ fn append_multiple_can_read_from_stdin() { )); assert_eq!( - assert_success(kvp(&with_dir(&dir, &["dump"]))), + assert_success(kvp(&with_dir(&dir, &["dump", "--text"]))), "x=1\nx=2\ny=3\n" ); } @@ -182,7 +199,7 @@ fn append_multiple_can_read_from_file() { &["append-multiple", "--file", input.to_str().unwrap()], ))); assert_eq!( - assert_success(kvp(&with_dir(&dir, &["dump"]))), + assert_success(kvp(&with_dir(&dir, &["dump", "--text"]))), "a=1\nb=2\n" ); } @@ -236,9 +253,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 +312,417 @@ fn report_failure_rejects_invalid_supporting_data() { .unwrap() .contains("key=value")); } + +#[test] +fn parsed_dump_reassembles_and_filters_without_dropping_other_entries() { + let dir = TempDir::new().unwrap(); + let base = format!( + "DIAG|agent|{VM_ID}|event|a:b|e5f01809-a7a3-4279-aa64-1f18e21eda6e|2026-08-31T00:00:00.000Z|none||" + ); + for (key, value) in [ + (format!("{base}|1"), "two"), + ("note".into(), "raw value"), + (format!("{base}|0"), "one/"), + ("DIAG_V2|future".into(), "preserved"), + ] { + assert_success(kvp(&with_dir( + &dir, + &["write", "--append", &key, value], + ))); + } + assert_success(kvp(&with_dir(&dir, &["report-success", "--vm-id", VM_ID]))); + assert_success(kvp(&with_dir( + &dir, + &[ + "emit", + "--name", + "ssh:key", + "--message", + "added", + "--vm-id", + VM_ID, + ], + ))); + let path = store_at(&dir).path().to_path_buf(); + let before = fs::read(&path).unwrap(); + + let entries = assert_json(kvp(&with_dir(&dir, &["dump", "--parse"]))); + assert_eq!(entries.as_array().unwrap().len(), 5); + assert_eq!(entries[0]["type"], "diagnostic"); + assert_eq!(entries[0]["kind"], "event"); + assert_eq!(entries[0]["name"], "a:b"); + assert_eq!(entries[0]["payload"], "one/two"); + assert_eq!( + entries[1], + json!({"type": "raw", "key": "note", "value": "raw value"}) + ); + assert_eq!( + entries[2], + json!({ + "type": "raw", "key": "DIAG_V2|future", "value": "preserved", + "error": "unsupported_version", + }) + ); + assert_eq!(entries[3]["type"], "PROVISIONING_REPORT"); + assert_eq!(entries[3]["result"], "success"); + assert_eq!(entries[4]["type"], "diagnostic"); + assert_eq!(entries[4]["name"], "ssh:key"); + + let filtered = assert_json(kvp(&with_dir( + &dir, + &["dump", "--parse", "--name", "ssh"], + ))); + assert_eq!( + filtered, + Value::Array(entries.as_array().unwrap()[1..].to_vec()) + ); + assert_eq!(fs::read(path).unwrap(), before); +} + +#[test] +fn parsed_dump_preserves_reader_pool_order() { + let dir = TempDir::new().unwrap(); + let store = store_at(&dir); + let event_id = "e5f01809-a7a3-4279-aa64-1f18e21eda6e"; + let key = |name: &str, timestamp: &str| { + format!( + "DIAG|agent|{VM_ID}|event|{name}|{event_id}|{timestamp}|none|||0" + ) + }; + // Timestamps are deliberately out of order to prove the CLI does not sort. + let records = [ + ("note".into(), "raw first"), + (key("latest", "2026-08-31T00:00:03.000Z"), "latest"), + ( + PROVISIONING_REPORT_KEY.into(), + "result=success|agent=agent|vm_id=vm|pps_type=None|timestamp=2026-08-31T00:00:02Z", + ), + ("DIAG_V2|future".into(), "raw second"), + (key("earliest", "2026-08-31T00:00:01.000Z"), "earliest"), + ]; + store + .append_multiple(records.iter().map(|(key, value)| (key, *value))) + .unwrap(); + let before = fs::read(store.path()).unwrap(); + + // Parsed output matches the reader entries verbatim (first-seen pool order). + let reader_entries = serde_json::to_value( + libazureinit_kvp::DiagnosticReader::new(store.clone()) + .entries() + .unwrap(), + ) + .unwrap(); + let parsed = assert_json(kvp(&with_dir(&dir, &["dump", "--parse"]))); + assert_eq!(parsed, reader_entries); + + // Physical dump keeps every record in file order. + let physical = assert_json(kvp(&with_dir(&dir, &["dump"]))); + let expected_physical: Vec<_> = records + .iter() + .map(|(key, value)| json!({"key": key, "value": value})) + .collect(); + assert_eq!(physical, Value::Array(expected_physical)); + assert_eq!(fs::read(store.path()).unwrap(), before); +} + +#[test] +fn parsed_dump_normalizes_cloud_init_in_json_and_text() { + 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 entries = assert_json(kvp(&with_dir(&dir, &["dump", "--parse"]))); + assert_eq!( + entries[0], + json!({ + "type": "diagnostic", "kind": "finish", "agent": "CLOUD_INIT", + "name": "modules-final/config-scripts_user", "vm_id": VM_ID, + "event_id": "e5f01809-a7a3-4279-aa64-1f18e21eda6e", + "timestamp": "2026-07-27T21:33:24.339006Z", "encoding": "none", + "result": "success", "duration": 0.5, "payload": "scripts ran", + }) + ); + assert_eq!(entries[1]["kind"], "start"); + assert!(entries[1].get("result").is_none()); + assert!(entries[1].get("duration").is_none()); + + let out = + assert_success(kvp(&with_dir(&dir, &["dump", "--parse", "--text"]))); + assert!(out.contains("diagnostic kind=finish")); + assert!(out.contains("agent=CLOUD_INIT")); + assert!(!out.contains("boot_epoch")); + 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.339006Z")); + assert!(out.contains("duration=0.5s")); + assert!(out.contains("payload=scripts ran")); + let start = out.lines().nth(1).unwrap(); + assert!(start.contains("diagnostic kind=start")); + assert!(!start.contains("result=")); + assert!(!start.contains("duration=")); +} + +#[rstest] +#[case(Encoding::GzB64, "gz+b64")] +#[case(Encoding::ZlibB64, "zlib+b64")] +fn parsed_dump_renders_bytes_reports_and_raw_errors( + #[case] encoding: Encoding, + #[case] token: &str, +) { + let dir = TempDir::new().unwrap(); + let store = store_at(&dir); + DiagnosticWriter::new(store.clone(), "agent", VM_ID) + .unwrap() + .emit_event( + "artifact", + vec![0, 255], + Some(encoding), + Some(Outcome::Failure), + Some(Duration::from_micros(7)), + ) + .unwrap(); + store.append("note", "raw value").unwrap(); + store.append("DIAG|bad", "junk").unwrap(); + assert_success(kvp(&with_dir( + &dir, + &["report-failure", "--vm-id", VM_ID, "--reason", "bad input"], + ))); + let out = + assert_success(kvp(&with_dir(&dir, &["dump", "--parse", "--text"]))); + let lines: Vec<_> = out.lines().collect(); + assert_eq!(lines.len(), 4); + assert!(lines[0].contains(&format!( + "encoding={token} result=fail duration=0.000007s payload_b64=AP8=" + ))); + assert_eq!(lines[1], "raw key=note value=raw value"); + assert_eq!(lines[2], "raw key=DIAG|bad value=junk error=malformed diagnostic or provisioning report"); + assert_eq!( + lines[3], + format!( + "PROVISIONING_REPORT={}", + store.read(PROVISIONING_REPORT_KEY).unwrap().unwrap() + ) + ); + + let entries = + assert_json(kvp(&with_dir(&dir, &["dump", "--parse", "--json"]))); + assert_eq!( + entries[0]["payload"], + json!({"type": "bytes", "encoding": "base64", "data": "AP8="}) + ); + assert_eq!(entries[2]["error"], "malformed"); + assert_eq!(entries[3]["reason"], "bad input"); +} + +#[test] +fn dump_name_requires_parse() { + let output = kvp(&["dump", "--name", "ssh"]); + assert_eq!(output.status.code(), Some(2)); + assert!(String::from_utf8(output.stderr) + .unwrap() + .contains("--parse")); +} + +#[test] +fn parsed_dump_filters_by_kind() { + let dir = TempDir::new().unwrap(); + let event_id = "e5f01809-a7a3-4279-aa64-1f18e21eda6e"; + let ts = "2026-08-31T00:00:00.000Z"; + let diag = |kind: &str, name: &str, result: &str, duration: &str| { + format!( + "DIAG|agent|{VM_ID}|{kind}|{name}|{event_id}|{ts}|none|{result}|{duration}|0" + ) + }; + for (key, value) in [ + (diag("start", "provision:run", "", ""), "starting"), + (diag("finish", "provision:run", "success", "312"), "done"), + (diag("event", "imds", "", ""), "ok"), + ("note".to_string(), "raw value"), + ] { + assert_success(kvp(&with_dir( + &dir, + &["write", "--append", &key, value], + ))); + } + assert_success(kvp(&with_dir(&dir, &["report-success", "--vm-id", VM_ID]))); + + // --kind keeps only diagnostics of that kind; reports and raw remain. + let finish = assert_json(kvp(&with_dir( + &dir, + &["dump", "--parse", "--kind", "finish"], + ))); + let finish = finish.as_array().unwrap(); + assert_eq!(finish.len(), 3); + assert_eq!(finish[0]["kind"], "finish"); + assert_eq!(finish[0]["name"], "provision:run"); + assert_eq!( + finish[1], + json!({"type": "raw", "key": "note", "value": "raw value"}) + ); + assert_eq!(finish[2]["type"], "PROVISIONING_REPORT"); + + // --name and --kind combine with AND semantics. + let combined = assert_json(kvp(&with_dir( + &dir, + &["dump", "--parse", "--kind", "start", "--name", "provision"], + ))); + let combined = combined.as_array().unwrap(); + assert_eq!(combined.len(), 3); + assert_eq!(combined[0]["kind"], "start"); + + // A kind/name pair matching no diagnostic keeps only reports and raw. + let empty = assert_json(kvp(&with_dir( + &dir, + &["dump", "--parse", "--kind", "start", "--name", "imds"], + ))); + let empty = empty.as_array().unwrap(); + assert_eq!(empty.len(), 2); + assert!(empty.iter().all(|entry| entry["type"] != "diagnostic")); +} + +#[test] +fn dump_kind_requires_parse_and_rejects_unknown_value() { + let requires_parse = kvp(&["dump", "--kind", "finish"]); + assert_eq!(requires_parse.status.code(), Some(2)); + assert!(String::from_utf8(requires_parse.stderr) + .unwrap() + .contains("--parse")); + + let unknown = kvp(&["dump", "--parse", "--kind", "bogus"]); + assert_eq!(unknown.status.code(), Some(2)); +} + +#[test] +fn conflicting_output_flags_fail_before_pool_access() { + let dir = TempDir::new().unwrap(); + let output = kvp(&with_dir(&dir, &["--json", "dump", "--text"])); + assert_eq!(output.status.code(), Some(2)); + assert!(output.stdout.is_empty()); + assert!(String::from_utf8(output.stderr) + .unwrap() + .contains("--json and --text cannot be used together")); + assert!(!store_at(&dir).path().exists()); +} + +#[test] +fn removed_diagnostic_options_are_rejected() { + let cases: &[(&[&str], &str)] = &[ + (&["dump", "--parse-diagnostics"], "--parse-diagnostics"), + (&["dump", "--parse", "--tail"], "--tail"), + (&["dump", "--parse", "-n", "1"], "-n"), + (&["emit", "--prefix", "agent"], "--prefix"), + ]; + for (args, flag) in cases { + let output = kvp(args); + assert_eq!(output.status.code(), Some(2)); + assert!(String::from_utf8(output.stderr) + .unwrap() + .contains(&format!("unexpected argument '{flag}'"))); + } +} + +#[test] +fn dumps_fail_without_partial_output_for_invalid_physical_utf8() { + let dir = TempDir::new().unwrap(); + let store = store_at(&dir); + store + .append_multiple([("good", "ok"), ("bad", "value")]) + .unwrap(); + let mut bytes = fs::read(store.path()).unwrap(); + let record_size = bytes.len() / 2; + bytes[record_size] = 0xff; + fs::write(store.path(), &bytes).unwrap(); + + for args in [ + &["dump"][..], + &["dump", "--parse"], + &["dump", "--parse", "--text"], + ] { + let output = kvp(&with_dir(&dir, args)); + assert_eq!(output.status.code(), Some(3)); + assert!(output.stdout.is_empty()); + assert!(!output.stderr.is_empty()); + } + assert_eq!(fs::read(store.path()).unwrap(), bytes); +} + +#[rstest] +#[case::default_agent(None)] +#[case::custom_agent(Some("azure-init-test"))] +fn emit_writes_event_readable_by_dump(#[case] agent: Option<&str>) { + let dir = TempDir::new().unwrap(); + let mut args = vec![ + "emit", + "--name", + "user:create_user", + "--message", + "created azureuser", + "--vm-id", + VM_ID, + ]; + if let Some(agent) = agent { + args.extend(["--agent", agent]); + } + assert_success(kvp(&with_dir(&dir, &args))); + + let entries = assert_json(kvp(&with_dir(&dir, &["dump", "--parse"]))); + let expected_agent = agent + .unwrap_or(concat!("libazureinit-kvp/", env!("CARGO_PKG_VERSION"))); + assert_eq!(entries.as_array().unwrap().len(), 1); + assert_eq!(entries[0]["type"], "diagnostic"); + assert_eq!(entries[0]["kind"], "event"); + assert_eq!(entries[0]["agent"], expected_agent); + assert_eq!(entries[0]["vm_id"], VM_ID); + assert_eq!(entries[0]["name"], "user:create_user"); + assert_eq!(entries[0]["payload"], "created azureuser"); + assert_eq!(entries[0]["encoding"], "none"); + assert!(entries[0].get("result").is_none()); + assert!(entries[0].get("duration").is_none()); + let event_id = + uuid::Uuid::parse_str(entries[0]["event_id"].as_str().unwrap()) + .unwrap(); + assert_eq!(event_id.get_version_num(), 4); + + let raw = assert_json(kvp(&with_dir(&dir, &["dump"]))); + let key = raw[0]["key"].as_str().unwrap(); + assert!(key.starts_with(&format!("DIAG|{expected_agent}|{VM_ID}|event|"))); + assert!(key.ends_with("|none|||0")); +} + +#[test] +fn emit_rejects_invalid_uuid_without_creating_pool() { + let dir = TempDir::new().unwrap(); + let output = kvp(&with_dir( + &dir, + &[ + "emit", + "--name", + "event", + "--message", + "test", + "--vm-id", + "vm-emit", + ], + )); + assert_eq!(output.status.code(), Some(2)); + assert!(output.stdout.is_empty()); + assert!(String::from_utf8(output.stderr).unwrap().contains("UUID")); + assert!(!store_at(&dir).path().exists()); +} diff --git a/libazureinit-kvp/tests/diagnostics.rs b/libazureinit-kvp/tests/diagnostics.rs new file mode 100644 index 00000000..8adbfee6 --- /dev/null +++ b/libazureinit-kvp/tests/diagnostics.rs @@ -0,0 +1,468 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +//! Public diagnostics API round trips, cloud-init compatibility, and store +//! behavior. + +use std::io::Read; +use std::thread; +use std::time::Duration; + +use base64::{engine::general_purpose::STANDARD, Engine as _}; +use flate2::read::ZlibDecoder; +use libazureinit_kvp::{ + write_report, DecodeError, Diagnostic, DiagnosticPayload, DiagnosticReader, + DiagnosticWriter, Encoding, Entry, Kind, KvpError, KvpPool, KvpPoolStore, + Outcome, PoolMode, ProvisioningReport, RawKeyValue, ReportPpsType, + DIAGNOSTIC_VERSION_ID, MAX_CHUNK_BYTES, +}; +use rstest::rstest; +use tempfile::TempDir; + +#[path = "fixtures/cloud_init.rs"] +mod cloud_init_fixtures; +use cloud_init_fixtures::COMPRESSED_LOG_CHUNKS; + +const AGENT: &str = "azure-init-test"; +const VM_ID: &str = "3f2504e0-4f89-41d3-9a0c-0305e82c3301"; +const EVENT_ID: &str = "8f3e9c4a-1b2c-4d5e-9f01-234567890abc"; + +fn store_at(dir: &TempDir) -> KvpPoolStore { + KvpPoolStore::new_in(KvpPool::Guest, dir.path(), PoolMode::Safe).unwrap() +} + +fn diagnostic(entry: &Entry) -> &Diagnostic { + let Entry::Diagnostic(diagnostic) = entry else { + panic!("expected a diagnostic, got {entry:?}"); + }; + diagnostic +} + +/// 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"}"#, + ), +]; + +#[rstest] +#[case::legacy(false)] +#[case::current(true)] +fn reads_real_cloud_init_pool_in_both_layouts(#[case] include_vm_id: bool) { + let dir = TempDir::new().unwrap(); + let store = store_at(&dir); + store + .append_multiple(CLOUD_INIT_RECORDS.iter().map(|&(key, value)| { + let key = if include_vm_id { + key.to_owned() + } else { + without_vm_id(key) + }; + (key, value) + })) + .unwrap(); + + let entries = DiagnosticReader::new(store).entries().unwrap(); + assert_eq!(entries.len(), CLOUD_INIT_RECORDS.len()); + + let Diagnostic::Finish(finish) = diagnostic(&entries[0]) else { + panic!("expected a finish"); + }; + assert_eq!(finish.key.agent, "CLOUD_INIT"); + assert_eq!(finish.key.name, "modules-final/config-scripts_user"); + assert_eq!( + finish.key.vm_id.as_deref(), + include_vm_id.then_some(CLOUD_INIT_VM_ID) + ); + assert_eq!(finish.result, Outcome::Success); + assert_eq!(finish.duration, Duration::from_micros(645)); + assert_eq!( + finish.payload, + DiagnosticPayload::from( + "config-scripts_user ran successfully and took 0.001 seconds" + ) + ); + assert!(matches!(diagnostic(&entries[1]), Diagnostic::Start(_))); + assert!(matches!(diagnostic(&entries[2]), Diagnostic::Finish(finish) + if finish.result == Outcome::Success + && finish.duration == Duration::from_micros(340712))); +} + +#[test] +fn span_and_point_events_round_trip_with_a_report() { + let dir = TempDir::new().unwrap(); + let store = store_at(&dir); + let writer = DiagnosticWriter::new(store.clone(), AGENT, VM_ID).unwrap(); + let reader = DiagnosticReader::new(store.clone()); + assert!(!store.path().exists()); + assert!(reader.entries().unwrap().is_empty()); + + writer + .emit_start(EVENT_ID, "provision:run", "starting", None) + .unwrap(); + writer + .emit_event( + "imds", + "ok", + None, + Some(Outcome::Success), + Some(Duration::from_micros(17)), + ) + .unwrap(); + writer + .emit_finish( + EVENT_ID, + "provision:run", + "finished", + None, + Outcome::Success, + Duration::from_micros(120), + ) + .unwrap(); + let report = ProvisioningReport::success(AGENT, VM_ID, ReportPpsType::None) + .with_extra("build", "test-123"); + write_report(&store, &report).unwrap(); + + let entries = reader.entries().unwrap(); + let [Entry::Diagnostic(Diagnostic::Start(start)), Entry::Diagnostic(Diagnostic::Event(event)), Entry::Diagnostic(Diagnostic::Finish(finish)), Entry::Report(decoded_report)] = + entries.as_slice() + else { + panic!("unexpected entries: {entries:?}"); + }; + assert_eq!(start.key.event_id, EVENT_ID); + assert_eq!(finish.key.event_id, EVENT_ID); + assert_eq!(start.key.name, finish.key.name); + assert_eq!(start.key.agent, AGENT); + assert_eq!(start.key.vm_id.as_deref(), Some(VM_ID)); + assert_eq!(start.payload, DiagnosticPayload::from("starting")); + assert_eq!(event.key.name, "imds"); + assert_eq!(event.payload, DiagnosticPayload::from("ok")); + assert_eq!(event.result, Some(Outcome::Success)); + assert_eq!(event.duration, Some(Duration::from_micros(17))); + assert_eq!( + uuid::Uuid::parse_str(&event.key.event_id) + .unwrap() + .get_version_num(), + 4 + ); + assert_ne!(event.key.event_id, EVENT_ID); + assert_eq!(finish.payload, DiagnosticPayload::from("finished")); + assert_eq!(finish.result, Outcome::Success); + assert_eq!(finish.duration, Duration::from_micros(120)); + assert_eq!(decoded_report, &report); + + let dumped = store.dump().unwrap(); + for (key, _) in &dumped[..3] { + assert!(key + .starts_with(&format!("{DIAGNOSTIC_VERSION_ID}|{AGENT}|{VM_ID}|"))); + assert!(key.ends_with("|0")); + } +} + +#[rstest] +#[case::text(None)] +#[case::gzip(Some(Encoding::GzB64))] +#[case::zlib(Some(Encoding::ZlibB64))] +fn long_payload_round_trips_through_host_visible_records( + #[case] encoding: Option, +) { + let dir = TempDir::new().unwrap(); + let store = store_at(&dir); + let writer = DiagnosticWriter::new(store.clone(), AGENT, VM_ID).unwrap(); + let message: String = (0..MAX_CHUNK_BYTES) + .map(|index| format!("{index:04x}\u{20ac};")) + .collect(); + writer + .emit_event( + "config:dump", + message.as_str(), + encoding.clone(), + None, + None, + ) + .unwrap(); + + let dumped = store.dump().unwrap(); + assert!(dumped.len() > 1); + assert_eq!(store.entries().unwrap().len(), dumped.len()); + let base = dumped[0].0.rsplit_once('|').unwrap().0; + for (index, (key, value)) in dumped.iter().enumerate() { + assert_eq!(key, &format!("{base}|{index}")); + assert!(key.len() <= 254); + assert!(value.len() <= MAX_CHUNK_BYTES); + } + + let decoded = + decode_single(DiagnosticReader::new(store).entries().unwrap()); + assert_eq!(decoded.kind(), Kind::Event); + assert_eq!(decoded.key().encoding, encoding); + let expected = if encoding.is_some() { + DiagnosticPayload::Bytes(message.into_bytes()) + } else { + DiagnosticPayload::Text(message) + }; + assert_eq!(decoded.payload(), &expected); +} + +#[test] +fn raw_and_malformed_records_are_preserved_beside_diagnostics() { + let dir = TempDir::new().unwrap(); + let store = store_at(&dir); + let records = vec![ + (format!("{AGENT}|100|{VM_ID}|event|legacy|{EVENT_ID}|2026-08-31T00:00:00Z|0"), "legacy", None), + ("DIAG|bad".into(), "junk", Some(DecodeError::Malformed)), + (format!("CLOUD_INIT|100|event|broken|{EVENT_ID}"), "not-json", Some(DecodeError::Malformed)), + ("PROVISIONING_REPORT".into(), "result=success", Some(DecodeError::Malformed)), + ("DIAG_V2|future".into(), "unknown", Some(DecodeError::UnsupportedVersion)), + ]; + store + .append_multiple(records.iter().map(|(key, value, _)| (key, *value))) + .unwrap(); + DiagnosticWriter::new(store.clone(), AGENT, VM_ID) + .unwrap() + .emit_event("valid", "visible", None, None, None) + .unwrap(); + let before = std::fs::read(store.path()).unwrap(); + + let entries = DiagnosticReader::new(store.clone()).entries().unwrap(); + assert_eq!(entries.len(), records.len() + 1); + for (entry, (key, value, error)) in entries.iter().zip(&records) { + assert_eq!( + entry, + &Entry::Raw(RawKeyValue { + key: key.clone(), + value: (*value).to_owned(), + error: *error, + }) + ); + } + assert_eq!( + diagnostic(entries.last().unwrap()).payload(), + &DiagnosticPayload::from("visible") + ); + assert_eq!(std::fs::read(store.path()).unwrap(), before); +} + +#[test] +fn chunked_entries_survive_store_swap_deletion() { + let dir = TempDir::new().unwrap(); + let store = store_at(&dir); + let writer = DiagnosticWriter::new(store.clone(), AGENT, VM_ID).unwrap(); + let first_message = "a".repeat(MAX_CHUNK_BYTES * 2 + 7); + let second_message = "b".repeat(MAX_CHUNK_BYTES * 2 + 7); + + store.append("remove-me", "raw").unwrap(); + writer + .emit_event("first", first_message.as_str(), None, None, None) + .unwrap(); + writer + .emit_event("second", second_message.as_str(), None, None, None) + .unwrap(); + + assert!(store.delete("remove-me").unwrap()); + + let entries = DiagnosticReader::new(store).entries().unwrap(); + assert_eq!(entries.len(), 2); + assert_eq!(diagnostic(&entries[0]).key().name, "second"); + assert_eq!( + diagnostic(&entries[0]).payload(), + &DiagnosticPayload::Text(second_message) + ); + assert_eq!(diagnostic(&entries[1]).key().name, "first"); + assert_eq!( + diagnostic(&entries[1]).payload(), + &DiagnosticPayload::Text(first_message) + ); +} + +#[test] +fn emit_rejects_delimiter_in_event_fields() { + let dir = TempDir::new().unwrap(); + let store = store_at(&dir); + let writer = DiagnosticWriter::new(store.clone(), AGENT, VM_ID).unwrap(); + + assert!(matches!( + writer.emit_event("a|b", "msg", None, None, None), + Err(KvpError::EventFieldContainsDelimiter { field: "name" }) + )); + assert!(!store.path().exists()); +} + +#[test] +fn concurrent_multichunk_emits_reassemble_without_interleaving() { + let dir = TempDir::new().unwrap(); + let store = store_at(&dir); + let writer = DiagnosticWriter::new(store.clone(), AGENT, VM_ID).unwrap(); + + const THREADS: usize = 5; + const PER_THREAD: usize = 8; + let len = MAX_CHUNK_BYTES * 2 + 7; + + let handles: Vec<_> = (0..THREADS) + .map(|thread_index| { + let writer = writer.clone(); + let marker = (b'a' + u8::try_from(thread_index).unwrap()) as char; + thread::spawn(move || { + for _ in 0..PER_THREAD { + let message = marker.to_string().repeat(len); + writer + .emit_event( + &format!("thread:{marker}"), + message, + None, + None, + None, + ) + .unwrap(); + } + }) + }) + .collect(); + for handle in handles { + handle.join().unwrap(); + } + + let dumped = store.dump().unwrap(); + assert_eq!(dumped.len(), THREADS * PER_THREAD * 3); + for group in dumped.chunks_exact(3) { + let base = group[0].0.rsplit_once('|').unwrap().0; + for (index, (key, _)) in group.iter().enumerate() { + assert_eq!(key, &format!("{base}|{index}")); + } + } + + let events = DiagnosticReader::new(store).entries().unwrap(); + assert_eq!(events.len(), THREADS * PER_THREAD); + for entry in &events { + let event = diagnostic(entry); + let DiagnosticPayload::Text(message) = event.payload() else { + panic!("expected text payload"); + }; + assert_eq!(message.len(), len); + let first = message.chars().next().unwrap(); + assert!(message.chars().all(|character| character == first)); + assert_eq!(event.key().name, format!("thread:{first}")); + } +} + +/// The reassembled `msg` across all three chunks. +const EXPECTED_COMPRESSED_MSG: &str = "{\"encoding\": \"gz+b64\", \"data\": \"eJxtV0uObDUMnbOKt4GHYjuJbYYIxAoYIQZx7EhIDBjAgN1ziumz1K2uqtY9lWuf3+XB++uwr2N/\\n4fEDfsb48vXLP3//8ef3f/37208///jrL7//8OVq2K09s4zekFEn71Lyp2fiBZsU1wxZ5zv+BpAa\\nwIzY128A1jhGRMikPJSsPJym6x5ih681gNwA6jzvnvK4U+6dhTdqu/Z4OsaY0y9xutDOkeIZDaw0\\nsKTytj1N9x3zVFjMoKnFtGaZNDCzganSNHnj8Hg4Am1eK4fcY2Nn7djuLuvJ5Iq5zVP2PPdWA78a\\n+CcyRXBfig14WKnskaNU6UWe1cDsBiZ0LamKNLIJIPyJVDl3nrHyzdh7lvt42F4DqQ3kfPXY2e9h\\nMdUyfhqYQDCJvsJoa4VY3KG1j5F3E7UGeC0GE9Xm5p0X5EzeRxYdnG+BCrV0F7Ht7qTeANrZ59Sk\\new4YjeU6mIjxSQNAnUZoGU1i2eDc3Ya71KQpo7m8UwTRfc/3oUFKk4c9cPbMoho6nhn+iQ8r1iQn\\nrdfAdrrAeeja0gxPp9w3d0FwV+LiVy4YVwMUH9hTp17qVOGQpwHZDs60VwTfqINRsY17Q2l9pmcR\\nj+4+qztrp5LYy/KCfmfbXbX4rlBfonb1NEZArRaW1+XDdtRTL0Nm+9LQLaI15t2Yx/tY18WJa3oD\\n22mDRtLaV1QGtHAfdvCSyavRKHVKyLev7u34WrOSISCLQQRnPgJ0sBskTw7Ts+hutlOBqN85Rplz\\nXuWwj6iyaAhjjIL7P5qqfM9rhEWdDl48ObTXmF44zUznfd+kby/nTgVF70HnA0sQ0h1SL/zZilP1\\nsAm4C6UW7ASEz7iBEzbQnUL8DD83hUtijQtvWjYOAogZfCGPdZMbtXKni0Ovzj2g/Ra4G3zZDVx5\\niKDbQHQqsHvXgQWB4ONOz21IHWKu3QB0fL8Yx8sNLlgG3Wv78IvN1xdpo0Xu+L6PrwejqVFX59R1\\nPgu86zGvSyDdXgQhra3I8Fmy2btddpwvQRbV2dc0RMaathBMrJCmyHIoynHqkQzv9VUb6x7Yazbw\\nnSIO3HFlbnWOElSMkncWwntin5KO1C5enuCOn6PdVDtN6COky/gkbjgh6su2R/qD/zFfHU0ccqcE\\nZKqmvPWK4Re+IovPxHi9tic+hU0DMO6zSfSpH9zcurRZYR+/NOS9ZwVxIRWLHzoFLB/rM4w469Nn\\nZoxN6ZMzpzdZIp1SRsbbG3Vj7YPocHpzudkUxxjxNQZCbIMxLKfAC1AuF0ljZdJpR8ZF3RvTEh0r\\nWN8q0Tu6y9tOlSOU5crBpfYxKA6y4ubyTjV8PsG639xy4Smw0oKp66B4oUdjolrFJ48uBFGjsVLp\\ndHQwaJ+JsISJokGIvIc6YBAkyLMPPIsnVo6CBKYf7QxVOhU5PDqPwe8hbFsuA2VSRx4nytmAtFqB\\n3SEDJbFIWM1aMI6gj82CNzJmV2ilU8dZajljvmFgSJrivtEeAaAo4rhlR/k5DLo1gK1KtvO8+TJj\\nfnxHH1wRwvj28tkp4RmthXGj7zAhqdCpBc8Ad+OMoIgQyTiINgGRdeNkzT5np4GHkNCZaFDCmDOC\\ntpYzTKFggG6JlhHPPiUykVQNaMf86yAGcgjJUQPlZI3/4xGPGPdsjsazZ6eAWHiUmugR8JI9HKx6\\n6NkPhfPDPopEZ1eFMTZbnZ0m8IQDlaPl/S8IG4GF8MJqcOD/AFeindw=\\n\"}"; + +const CLOUD_INIT_VM_ID: &str = "0e5e179d-5341-478b-8456-fbb90621bdf8"; + +/// Insert a `vm_id` segment after `name`, turning an old-format key into +/// the current layout (any trailing chunk index is preserved): +/// `CLOUD_INIT|inc|type|name|uuid[|i]` +/// -> `CLOUD_INIT|inc|type|name|vm_id|uuid[|i]`. +fn with_vm_id(old_key: &str, vm_id: &str) -> String { + let mut segments: Vec<&str> = old_key.split('|').collect(); + segments.insert(4, vm_id); + segments.join("|") +} + +fn without_vm_id(current_key: &str) -> String { + let mut segments: Vec<&str> = current_key.split('|').collect(); + segments.remove(4); + segments.join("|") +} + +/// Append the given records to a fresh guest pool and normalize them. +fn entries_of, V: AsRef>(pairs: &[(K, V)]) -> Vec { + let dir = TempDir::new().unwrap(); + let store = store_at(&dir); + store + .append_multiple( + pairs + .iter() + .map(|(key, value)| (key.as_ref(), value.as_ref())), + ) + .unwrap(); + DiagnosticReader::new(store).entries().unwrap() +} + +fn decode_single(entries: Vec) -> Diagnostic { + assert_eq!(entries.len(), 1, "expected one entry, got: {entries:?}"); + match entries.into_iter().next().unwrap() { + Entry::Diagnostic(diagnostic) => diagnostic, + entry => panic!("expected a diagnostic, got {entry:?}"), + } +} + +#[rstest] +#[case::legacy(false)] +#[case::current(true)] +fn captured_compressed_log_reassembles_and_decodes( + #[case] include_vm_id: bool, +) { + let records: Vec<(String, &str)> = COMPRESSED_LOG_CHUNKS + .iter() + .rev() + .map(|&(key, value)| { + let key = if include_vm_id { + with_vm_id(key, CLOUD_INIT_VM_ID) + } else { + key.to_owned() + }; + (key, value) + }) + .collect(); + let event = decode_single(entries_of(&records)); + assert_eq!(event.kind(), Kind::Event); + assert_eq!( + event.key().vm_id.as_deref(), + include_vm_id.then_some(CLOUD_INIT_VM_ID) + ); + assert_eq!(event.key().name, "cloud-init.log"); + assert_eq!(event.key().encoding, Some(Encoding::GzB64)); + + let envelope: serde_json::Value = + serde_json::from_str(EXPECTED_COMPRESSED_MSG).unwrap(); + let data: String = envelope["data"] + .as_str() + .unwrap() + .split_ascii_whitespace() + .collect(); + let compressed = STANDARD.decode(data).unwrap(); + let mut expected = Vec::new(); + ZlibDecoder::new(compressed.as_slice()) + .read_to_end(&mut expected) + .unwrap(); + assert_eq!(expected.len(), 3554); + assert_eq!(event.payload(), &DiagnosticPayload::Bytes(expected)); +} + +#[test] +fn incomplete_cloud_init_group_preserves_each_physical_record() { + let base = "CLOUD_INIT|1786047606|event|test|\ + b7a822ba-4eea-46c0-b559-e84396101132"; + let chunk = |index: u32, message: &str| { + format!( + r#"{{"name":"test","type":"event","ts":"2026-08-06T20:20:13Z","msg_i":{index},"msg":"{message}"}}"# + ) + }; + let records = vec![ + (format!("{base}|2"), chunk(2, "third")), + ("note".to_owned(), "untouched".to_owned()), + (format!("{base}|0"), chunk(0, "first")), + ]; + + let entries = entries_of(&records); + let expected: Vec<_> = records + .into_iter() + .map(|(key, value)| { + let error = (key != "note").then_some(DecodeError::IncompleteGroup); + Entry::Raw(RawKeyValue { key, value, error }) + }) + .collect(); + assert_eq!(entries, expected); +} diff --git a/libazureinit-kvp/tests/fixtures/cloud_init.rs b/libazureinit-kvp/tests/fixtures/cloud_init.rs new file mode 100644 index 00000000..3eafd5e3 --- /dev/null +++ b/libazureinit-kvp/tests/fixtures/cloud_init.rs @@ -0,0 +1,17 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +pub const COMPRESSED_LOG_CHUNKS: &[(&str, &str)] = &[ + ( + "CLOUD_INIT|1786047606|compressed|cloud-init.log|b7a822ba-4eea-46c0-b559-e84396101132|0", + "{\"name\":\"cloud-init.log\",\"type\":\"compressed\",\"ts\":\"2026-08-06T20:20:13.479078Z\",\"msg_i\":0,\"msg\":\"{\\\"encoding\\\": \\\"gz+b64\\\", \\\"data\\\": \\\"eJxtV0uObDUMnbOKt4GHYjuJbYYIxAoYIQZx7EhIDBjAgN1ziumz1K2uqtY9lWuf3+XB++uwr2N/\\\\n4fEDfsb48vXLP3//8ef3f/37208///jrL7//8OVq2K09s4zekFEn71Lyp2fiBZsU1wxZ5zv+BpAa\\\\nwIzY128A1jhGRMikPJSsPJym6x5ih681gNwA6jzvnvK4U+6dhTdqu/Z4OsaY0y9xutDOkeIZDaw0\\\\nsKTytj1N9x3zVFjMoKnFtGaZNDCzganSNHnj8Hg4Am1eK4fcY2Nn7djuLuvJ5Iq5zVP2PPdWA78a\\\\n+CcyRXBfig14WKnskaNU6UWe1cDsBiZ0LamKNLIJIPyJVDl3nrHyzdh7lvt42F4DqQ3kfPXY2e9h\\\\nMdUyfhqYQDCJvsJoa4VY3KG1j5F3E7UGeC0GE9Xm5p0X5EzeRxYdnG+BCrV0F7Ht7qTeANrZ59Sk\\\\new4YjeU6mIjxSQNAnUZoGU1i2eDc3Ya71KQpo7m8UwTRfc/3oUFKk4c9cPbMoho6nhn+iQ8r1iQn\\\\nrdfAdrrAeeja0gxPp9w3d0FwV+LiVy4YVwMUH9hTp17qVOGQpwHZDs60VwTfqINRsY17Q2l9pmcR\\\\nj+4+qztrp5LYy/KCfmfbXbX4rlBfonb1NEZArRaW1+XDdtRTL0Nm+9LQLaI15t2Yx/tY18WJa3oD\\\\n22mDRtLaV1QGtHAfdvCSyavRKHVKyLev7u34WrOSISCLQQRnPgJ0sBskTw7Ts+hutlOBqN85Rplz\\\\nXuWwj6iyaAhjjIL7P5qqfM9rhEWdDl48ObTXmF44zUznfd+kby/nTgVF70HnA0sQ0h1SL/zZilP1\\\\nsAm4C6UW7\"}", + ), + ( + "CLOUD_INIT|1786047606|compressed|cloud-init.log|b7a822ba-4eea-46c0-b559-e84396101132|1", + "{\"name\":\"cloud-init.log\",\"type\":\"compressed\",\"ts\":\"2026-08-06T20:20:13.479078Z\",\"msg_i\":1,\"msg\":\"ASEz7iBEzbQnUL8DD83hUtijQtvWjYOAogZfCGPdZMbtXKni0Ovzj2g/Ra4G3zZDVx5\\\\niKDbQHQqsHvXgQWB4ONOz21IHWKu3QB0fL8Yx8sNLlgG3Wv78IvN1xdpo0Xu+L6PrwejqVFX59R1\\\\nPgu86zGvSyDdXgQhra3I8Fmy2btddpwvQRbV2dc0RMaathBMrJCmyHIoynHqkQzv9VUb6x7Yazbw\\\\nnSIO3HFlbnWOElSMkncWwntin5KO1C5enuCOn6PdVDtN6COky/gkbjgh6su2R/qD/zFfHU0ccqcE\\\\nZKqmvPWK4Re+IovPxHi9tic+hU0DMO6zSfSpH9zcurRZYR+/NOS9ZwVxIRWLHzoFLB/rM4w469Nn\\\\nZoxN6ZMzpzdZIp1SRsbbG3Vj7YPocHpzudkUxxjxNQZCbIMxLKfAC1AuF0ljZdJpR8ZF3RvTEh0r\\\\nWN8q0Tu6y9tOlSOU5crBpfYxKA6y4ubyTjV8PsG639xy4Smw0oKp66B4oUdjolrFJ48uBFGjsVLp\\\\ndHQwaJ+JsISJokGIvIc6YBAkyLMPPIsnVo6CBKYf7QxVOhU5PDqPwe8hbFsuA2VSRx4nytmAtFqB\\\\n3SEDJbFIWM1aMI6gj82CNzJmV2ilU8dZajljvmFgSJrivtEeAaAo4rhlR/k5DLo1gK1KtvO8+TJj\\\\nfnxHH1wRwvj28tkp4RmthXGj7zAhqdCpBc8Ad+OMoIgQyTiINgGRdeNkzT5np4GHkNCZaFDCmDOC\\\\ntpYzTKFggG6JlhHPPiUykVQNaMf86yAGcgjJUQPlZI3/4xGPGPdsjsazZ6eAWHiUmugR8JI9HKx6\\\\n6NkPhfPDPopEZ1eFMTZbnZ0m8IQDlaPl/S8IG4GF8MJqcOD/AFeindw=\\\"}", + ), + ( + "CLOUD_INIT|1786047606|compressed|cloud-init.log|b7a822ba-4eea-46c0-b559-e84396101132|2", + "{\"name\":\"cloud-init.log\",\"type\":\"compressed\",\"ts\":\"2026-08-06T20:20:13.479078Z\",\"msg_i\":2,\"msg\":\"\\n\\\"}\"}", + ), +];