badpipe/src/main.rs

133 lines
3.7 KiB
Rust

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); }
}