Merge remote-tracking branch 'origin/main' into bruceg/OPA-4723-add-templated-json-generator

Author: bruceg

CHANGELOG.md

@@ -37,6 +37,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Added new payload generator `templated_json` that uses a user-specificied
   payload template to produce JSON lines. See
   `lading_payload/README.templated_json.md` for more details.
+- `payloadtool` now supports a `--dump=FILENAME` option which saves the generated data
+  blocks to the named file.
 ## Fixed
 - Fixed a race condition in `lading_signal` that caused lading to hang on shutdown.
 - Fixed a tag parsing bug that resulted in tags with hyphenated keys being merged with proceeding tag values.

lading/src/bin/payloadtool.rs

@@ -139,7 +139,7 @@ mod alloc_tracker {
 
 use std::fmt;
 use std::fs::{File, OpenOptions};
-use std::io::Read;
+use std::io::{Read, Write};
 use std::num::NonZeroU32;
 use std::path::{Path, PathBuf};
 use std::time::Instant;
@@ -150,7 +150,7 @@ use tokio::runtime::Builder;
 use anyhow::{Context, Result, anyhow};
 use byte_unit::{Byte, Unit, UnitType};
 use clap::Parser;
-use lading::generator::{self, http::Method};
+use lading::generator::{self, file_gen, http::Method};
 use lading_payload::block;
 use rand::{SeedableRng, rngs::StdRng};
 use sha2::{Digest, Sha256};
@@ -276,14 +276,17 @@ struct Args {
     /// Report memory allocation statistics at completion
     #[clap(short = 'm', long)]
     memory_stats: bool,
+    /// Dump generated payload data to this file
+    #[clap(long)]
+    dump: Option<PathBuf>,
 }
 
 fn generate_and_check(
     config: &lading_payload::Config,
     seed: [u8; 32],
     total_bytes: NonZeroU32,
     max_block_size: Byte,
-    compute_fingerprint: bool,
+    args: &Args,
 ) -> Result<Option<Fingerprint>> {
     let mut rng = StdRng::from_seed(seed);
     let start = Instant::now();
@@ -306,8 +309,23 @@ fn generate_and_check(
     info!("Payload generation took {:?}", start.elapsed());
     trace!("Payload: {:#?}", blocks);
 
+    if let Some(dump_path) = args.dump.as_deref() {
+        let mut dump_file = OpenOptions::new()
+            .write(true)
+            .create(true)
+            .truncate(true)
+            .open(dump_path)
+            .with_context(|| format!("Could not open dump file at: {}", dump_path.display()))?;
+        for block in &blocks {
+            dump_file.write_all(&block.bytes).with_context(|| {
+                format!("Failed to write to dump file: {}", dump_path.display())
+            })?;
+        }
+        info!("Dumped payload to {}", dump_path.display());
+    }
+
     // Compute fingerprint if requested: SHA256 hash and Shannon entropy.
-    let fingerprint = if compute_fingerprint {
+    let fingerprint = args.fingerprint.then(|| {
         let mut hasher = Sha256::new();
         let mut all_bytes = Vec::new();
         for block in &blocks {
@@ -317,10 +335,8 @@ fn generate_and_check(
         let result = hasher.finalize();
         let hash = format!("{result:x}");
         let entropy = shannon_entropy(&all_bytes);
-        Some(Fingerprint { hash, entropy })
-    } else {
-        None
-    };
+        Fingerprint { hash, entropy }
+    });
 
     let mut total_generated_bytes: u32 = 0;
     for block in &blocks {
@@ -348,55 +364,54 @@ fn generate_and_check(
 }
 
 #[expect(clippy::too_many_lines)]
-fn check_generator(
-    config: &generator::Config,
-    compute_fingerprint: bool,
-) -> Result<Option<Fingerprint>> {
+fn check_generator(config: &generator::Config, args: &Args) -> Result<Option<Fingerprint>> {
     match &config.inner {
-        generator::Inner::FileGen(_) => {
-            if compute_fingerprint {
-                warn!("FileGen not supported for fingerprinting");
-                return Ok(None);
-            }
-            unimplemented!("FileGen not supported")
+        generator::Inner::FileGen(g) => {
+            let (cache_size, variant, seed, max_block_size) = match g {
+                file_gen::Config::Traditional(c) => (
+                    c.maximum_prebuild_cache_size_bytes,
+                    &c.variant,
+                    c.seed,
+                    c.maximum_block_size,
+                ),
+                file_gen::Config::Logrotate(c) => (
+                    c.maximum_prebuild_cache_size_bytes,
+                    &c.variant,
+                    c.seed,
+                    c.maximum_block_size,
+                ),
+                #[cfg(feature = "logrotate_fs")]
+                file_gen::Config::LogrotateFs(c) => (
+                    c.maximum_prebuild_cache_size_bytes,
+                    &c.variant,
+                    c.seed,
+                    c.maximum_block_size,
+                ),
+            };
+            #[expect(clippy::cast_possible_truncation)]
+            let total_bytes = NonZeroU32::new(cache_size.as_u128() as u32)
+                .expect("Non-zero max prebuild cache size");
+            generate_and_check(variant, seed, total_bytes, max_block_size, args)
         }
         generator::Inner::UnixDatagram(g) => {
             let max_block_size = UDP_PACKET_LIMIT_BYTES;
             #[expect(clippy::cast_possible_truncation)]
             let total_bytes = NonZeroU32::new(g.maximum_prebuild_cache_size_bytes.as_u128() as u32)
                 .expect("Non-zero max prebuild cache size");
-            generate_and_check(
-                &g.variant,
-                g.seed,
-                total_bytes,
-                max_block_size,
-                compute_fingerprint,
-            )
+            generate_and_check(&g.variant, g.seed, total_bytes, max_block_size, args)
         }
         generator::Inner::Tcp(g) => {
             #[expect(clippy::cast_possible_truncation)]
             let total_bytes = NonZeroU32::new(g.maximum_prebuild_cache_size_bytes.as_u128() as u32)
                 .expect("Non-zero max prebuild cache size");
-            generate_and_check(
-                &g.variant,
-                g.seed,
-                total_bytes,
-                g.maximum_block_size,
-                compute_fingerprint,
-            )
+            generate_and_check(&g.variant, g.seed, total_bytes, g.maximum_block_size, args)
         }
         generator::Inner::Udp(g) => {
             #[expect(clippy::cast_possible_truncation)]
             let total_bytes = NonZeroU32::new(g.maximum_prebuild_cache_size_bytes.as_u128() as u32)
                 .expect("Non-zero max prebuild cache size");
             let max_block_size = UDP_PACKET_LIMIT_BYTES;
-            generate_and_check(
-                &g.variant,
-                g.seed,
-                total_bytes,
-                max_block_size,
-                compute_fingerprint,
-            )
+            generate_and_check(&g.variant, g.seed, total_bytes, max_block_size, args)
         }
         generator::Inner::Http(g) => {
             let (variant, max_prebuild_cache_size_bytes) = match &g.method {
@@ -409,23 +424,17 @@ fn check_generator(
             #[expect(clippy::cast_possible_truncation)]
             let total_bytes = NonZeroU32::new(max_prebuild_cache_size_bytes.as_u128() as u32)
                 .expect("Non-zero max prebuild cache size");
-            generate_and_check(
-                variant,
-                g.seed,
-                total_bytes,
-                g.maximum_block_size,
-                compute_fingerprint,
-            )
+            generate_and_check(variant, g.seed, total_bytes, g.maximum_block_size, args)
         }
         generator::Inner::SplunkHec(_) => {
-            if compute_fingerprint {
+            if args.fingerprint {
                 warn!("SplunkHec not supported for fingerprinting");
                 return Ok(None);
             }
             unimplemented!("SplunkHec not supported")
         }
         generator::Inner::FileTree(_) => {
-            if compute_fingerprint {
+            if args.fingerprint {
                 warn!("FileTree not supported for fingerprinting");
                 return Ok(None);
             }
@@ -435,61 +444,43 @@ fn check_generator(
             #[expect(clippy::cast_possible_truncation)]
             let total_bytes = NonZeroU32::new(g.maximum_prebuild_cache_size_bytes.as_u128() as u32)
                 .expect("Non-zero max prebuild cache size");
-            generate_and_check(
-                &g.variant,
-                g.seed,
-                total_bytes,
-                g.maximum_block_size,
-                compute_fingerprint,
-            )
+            generate_and_check(&g.variant, g.seed, total_bytes, g.maximum_block_size, args)
         }
         generator::Inner::UnixStream(g) => {
             #[expect(clippy::cast_possible_truncation)]
             let total_bytes = NonZeroU32::new(g.maximum_prebuild_cache_size_bytes.as_u128() as u32)
                 .expect("Non-zero max prebuild cache size");
-            generate_and_check(
-                &g.variant,
-                g.seed,
-                total_bytes,
-                g.maximum_block_size,
-                compute_fingerprint,
-            )
+            generate_and_check(&g.variant, g.seed, total_bytes, g.maximum_block_size, args)
         }
         generator::Inner::PassthruFile(g) => {
             #[expect(clippy::cast_possible_truncation)]
             let total_bytes = NonZeroU32::new(g.maximum_prebuild_cache_size_bytes.as_u128() as u32)
                 .expect("Non-zero max prebuild cache size");
-            generate_and_check(
-                &g.variant,
-                g.seed,
-                total_bytes,
-                g.maximum_block_size,
-                compute_fingerprint,
-            )
+            generate_and_check(&g.variant, g.seed, total_bytes, g.maximum_block_size, args)
         }
         generator::Inner::ProcessTree(_) => {
-            if compute_fingerprint {
+            if args.fingerprint {
                 warn!("ProcessTree not supported for fingerprinting");
                 return Ok(None);
             }
             unimplemented!("ProcessTree not supported")
         }
         generator::Inner::ProcFs(_) => {
-            if compute_fingerprint {
+            if args.fingerprint {
                 warn!("ProcFs not supported for fingerprinting");
                 return Ok(None);
             }
             unimplemented!("ProcFs not supported")
         }
         generator::Inner::Container(_) => {
-            if compute_fingerprint {
+            if args.fingerprint {
                 warn!("Container not supported for fingerprinting");
                 return Ok(None);
             }
             unimplemented!("Container not supported")
         }
         generator::Inner::Kubernetes(_) => {
-            if compute_fingerprint {
+            if args.fingerprint {
                 warn!("Kubernetes not supported for fingerprinting");
                 return Ok(None);
             }
@@ -500,13 +491,7 @@ fn check_generator(
                 generator::trace_agent::validate_cache_size(g.maximum_prebuild_cache_size_bytes)
                     .map_err(|e| anyhow::anyhow!("Cache size validation failed: {e}"))?;
             let conf = lading_payload::Config::TraceAgent(g.variant);
-            generate_and_check(
-                &conf,
-                g.seed,
-                total_bytes,
-                g.maximum_block_size,
-                compute_fingerprint,
-            )
+            generate_and_check(&conf, g.seed, total_bytes, g.maximum_block_size, args)
         }
     }
 }
@@ -548,24 +533,24 @@ async fn inner_main() -> Result<()> {
         config.generator.len()
     );
 
-    if let Some(generator_id) = args.generator_id {
+    if let Some(ref generator_id) = args.generator_id {
         let generator = config
             .generator
             .iter()
             .find(|g| {
                 let Some(ref id) = g.general.id else {
                     return false;
                 };
-                *id == generator_id
+                id == generator_id
             })
             .ok_or_else(|| anyhow!("No generator found with id: {generator_id}"))?;
-        let fingerprint = check_generator(generator, args.fingerprint)?;
+        let fingerprint = check_generator(generator, &args)?;
         if args.fingerprint
             && let Some(fp) = fingerprint
         {
-            if let Some(verify_path) = args.verify {
+            if let Some(ref verify_path) = args.verify {
                 let expected_content =
-                    fs::read_to_string(&verify_path).await.with_context(|| {
+                    fs::read_to_string(verify_path).await.with_context(|| {
                         format!("Could not read verify file {}", verify_path.display())
                     })?;
 
@@ -597,19 +582,17 @@ async fn inner_main() -> Result<()> {
     } else {
         let mut all_fingerprints = Vec::new();
         for generator in config.generator {
-            let fingerprint = check_generator(&generator, args.fingerprint)?;
-            if args.fingerprint
-                && let Some(fp) = fingerprint
-            {
+            let fingerprint = check_generator(&generator, &args)?;
+            if let Some(fp) = fingerprint {
                 let gen_id = generator.general.id.as_deref().unwrap_or("<unnamed>");
                 all_fingerprints.push((gen_id.to_string(), fp));
             }
         }
         if args.fingerprint && !all_fingerprints.is_empty() {
-            if let Some(verify_path) = args.verify {
+            if let Some(ref verify_path) = args.verify {
                 // Read expected fingerprints from file
                 let expected_content =
-                    fs::read_to_string(&verify_path).await.with_context(|| {
+                    fs::read_to_string(verify_path).await.with_context(|| {
                         format!("Could not read verify file {}", verify_path.display())
                     })?;
 

lading/src/generator/file_gen/logrotate.rs

@@ -140,10 +140,10 @@ pub struct Config {
     bytes_per_second: Option<Byte>,
     /// Defines the maximum internal cache of this log target. `file_gen` will
     /// pre-build its outputs up to the byte capacity specified here.
-    maximum_prebuild_cache_size_bytes: Byte,
+    pub maximum_prebuild_cache_size_bytes: Byte,
     /// The maximum size in bytes of the largest block in the prebuild cache.
     #[serde(default = "lading_payload::block::default_maximum_block_size")]
-    maximum_block_size: byte_unit::Byte,
+    pub maximum_block_size: byte_unit::Byte,
     /// Whether to use a fixed or streaming block cache
     #[serde(default = "lading_payload::block::default_cache_method")]
     block_cache_method: block::CacheMethod,

lading/src/generator/file_gen/logrotate_fs.rs

@@ -51,13 +51,13 @@ pub struct Config {
     /// files will be present in the root path.
     max_depth: u8,
     /// Sets the [`crate::payload::Config`] of this template.
-    variant: lading_payload::Config,
+    pub variant: lading_payload::Config,
     /// Defines the maximum internal cache of this log target. `file_gen` will
     /// pre-build its outputs up to the byte capacity specified here.
-    maximum_prebuild_cache_size_bytes: byte_unit::Byte,
+    pub maximum_prebuild_cache_size_bytes: byte_unit::Byte,
     /// The maximum size in bytes of the largest block in the prebuild cache.
     #[serde(default = "lading_payload::block::default_maximum_block_size")]
-    maximum_block_size: byte_unit::Byte,
+    pub maximum_block_size: byte_unit::Byte,
     /// The mount-point for this filesystem
     mount_point: PathBuf,
     /// The load profile, controlling bytes or blocks per second as a function of time.

lading/src/generator/file_gen/traditional.rs

@@ -107,10 +107,10 @@ pub struct Config {
     bytes_per_second: Option<Byte>,
     /// Defines the maximum internal cache of this log target. `file_gen` will
     /// pre-build its outputs up to the byte capacity specified here.
-    maximum_prebuild_cache_size_bytes: Byte,
+    pub maximum_prebuild_cache_size_bytes: Byte,
     /// The maximum size in bytes of the largest block in the prebuild cache.
     #[serde(default = "lading_payload::block::default_maximum_block_size")]
-    maximum_block_size: byte_unit::Byte,
+    pub maximum_block_size: byte_unit::Byte,
     /// Whether to use a fixed or streaming block cache
     #[serde(default = "lading_payload::block::default_cache_method")]
     block_cache_method: block::CacheMethod,