basic error-injecting pipe

src/main.rs

@@ -0,0 +1,132 @@
+use clap::{Parser, Args};
+use rand::{Rng, distributions::{Distribution, Uniform}};
+
+use std::io::{self, Read, Write};
+
+const BUF_SIZE: usize = 8192;
+
+#[derive(Parser, Debug)]
+#[clap(author, version, about)]
+struct Cli {
+    #[clap(short, long)]
+    verbose: bool,
+
+    /// Introduce an error every N bytes on average
+    #[clap(short, long, value_name = "N", default_value_t = 1000)]
+    rarity: u32,
+
+    #[clap(flatten, help_heading = "Types of errors to cause. If none are specified individually, all will be allowed")]
+    error_spec: ErrorSpec,
+}
+
+#[derive(Args, Debug, Clone, Copy)]
+struct ErrorSpec {
+    /// Insert extra bytes into the stream
+    #[clap(short, long)]
+    insert: bool,
+
+    /// Delete bytes from the stream
+    #[clap(short, long)]
+    delete: bool,
+
+    /// Alter bytes in the stream
+    #[clap(short, long)]
+    alter: bool,
+}
+
+enum Error {
+    Insert(u8),
+    Delete,
+    Alter(u8),
+}
+
+impl ErrorSpec {
+    fn gen(self) -> Error {
+        let spec = if let ErrorSpec { insert: false, delete: false, alter: false } = self {
+            ErrorSpec { insert: true, delete: true, alter: true }
+        } else {
+            self
+        };
+
+        let mut options_remaining =
+            if spec.insert { 1 } else { 0 } +
+            if spec.delete { 1 } else { 0 } +
+            if spec.alter { 1 } else { 0 };
+
+        let mut rng = rand::thread_rng();
+
+        if spec.insert {
+            if rng.gen_range(0..options_remaining) == 0 {
+                return Error::Insert(rng.gen())
+            }
+            options_remaining -= 1;
+        }
+        if spec.delete {
+            if rng.gen_range(0..options_remaining) == 0 {
+                return Error::Delete
+            }
+            options_remaining -= 1;
+        }
+        if spec.alter {
+            if rng.gen_range(0..options_remaining) == 0 {
+                return Error::Alter(rng.gen())
+            }
+            options_remaining -= 1;
+        }
+        panic!("Expected 0 options remaining, got {}, this shouldn't happen", options_remaining)
+    }
+}
+
+fn read_handling_interruptions<T: Read>(source: &mut T, dest: &mut [u8]) -> io::Result<usize> {
+    loop {
+        let result = source.read(dest);
+        if let Err(e) = &result {
+            if let io::ErrorKind::Interrupted = e.kind() {
+                continue
+            }
+        }
+        return result
+    }
+}
+
+fn main() {
+    let cli = Cli::parse();
+
+    let mut stdin = io::stdin();
+    let mut stdout = io::stdout();
+    let mut buffer = [0u8; BUF_SIZE];
+    let mut total = 0;
+    let rarity = Uniform::from(0..cli.rarity);
+    let mut rng = rand::thread_rng();
+
+    loop {
+        let fill = read_handling_interruptions(&mut stdin, &mut buffer).unwrap();
+        if fill == 0 {
+            break
+        }
+        let mut pos = 0;
+        let error_indices = (0..fill).filter(|_| rarity.sample(&mut rng) == 0);
+        for idx in error_indices {
+            stdout.write_all(&buffer[pos..idx]).unwrap();
+            if cli.verbose { eprint!("@ 0x{:08x}: ", total + idx); }
+            match cli.error_spec.gen() {
+                Error::Insert(b) => {
+                    stdout.write_all(&[buffer[idx], b]).unwrap();
+                    if cli.verbose { eprintln!("Inserted byte {:02x}", b); }
+                }
+                Error::Delete => {
+                    if cli.verbose { eprintln!("Deleted byte {:02x}", buffer[idx]); }
+                }
+                Error::Alter(b) => {
+                    stdout.write_all(&[b]).unwrap();
+                    if cli.verbose { eprintln!("Altered byte {:02x} to {:02x}", buffer[idx], b); }
+                }
+            }
+            pos = idx + 1;
+        }
+        stdout.write_all(&buffer[pos..fill]).unwrap();
+        total += fill;
+    }
+
+    if cli.verbose { eprintln!("Reached EOF after {} bytes", total); }
+}