Yecheng Zhao | fabac8e | 2024-02-17 23:18:12 +0000 | [diff] [blame] | 1 | #![allow(dead_code)] |
| 2 | |
| 3 | #[cfg(feature = "log")] |
| 4 | use env_logger::Builder; |
| 5 | use getopts::{Matches, Options}; |
| 6 | #[cfg(feature = "log")] |
| 7 | use log::{trace, Level, LevelFilter}; |
| 8 | use std::env; |
| 9 | use std::fs::File; |
| 10 | use std::io::{self, Write}; |
| 11 | use std::process; |
| 12 | use std::str::{self, FromStr}; |
| 13 | use std::time::{SystemTime, UNIX_EPOCH}; |
| 14 | |
| 15 | #[cfg(feature = "phy-tuntap_interface")] |
| 16 | use smoltcp::phy::TunTapInterface; |
| 17 | use smoltcp::phy::{Device, FaultInjector, Medium, Tracer}; |
| 18 | use smoltcp::phy::{PcapMode, PcapWriter}; |
| 19 | use smoltcp::time::{Duration, Instant}; |
| 20 | |
| 21 | #[cfg(feature = "log")] |
| 22 | pub fn setup_logging_with_clock<F>(filter: &str, since_startup: F) |
| 23 | where |
| 24 | F: Fn() -> Instant + Send + Sync + 'static, |
| 25 | { |
| 26 | Builder::new() |
| 27 | .format(move |buf, record| { |
| 28 | let elapsed = since_startup(); |
| 29 | let timestamp = format!("[{elapsed}]"); |
| 30 | if record.target().starts_with("smoltcp::") { |
| 31 | writeln!( |
| 32 | buf, |
| 33 | "\x1b[0m{} ({}): {}\x1b[0m", |
| 34 | timestamp, |
| 35 | record.target().replace("smoltcp::", ""), |
| 36 | record.args() |
| 37 | ) |
| 38 | } else if record.level() == Level::Trace { |
| 39 | let message = format!("{}", record.args()); |
| 40 | writeln!( |
| 41 | buf, |
| 42 | "\x1b[37m{} {}\x1b[0m", |
| 43 | timestamp, |
| 44 | message.replace('\n', "\n ") |
| 45 | ) |
| 46 | } else { |
| 47 | writeln!( |
| 48 | buf, |
| 49 | "\x1b[32m{} ({}): {}\x1b[0m", |
| 50 | timestamp, |
| 51 | record.target(), |
| 52 | record.args() |
| 53 | ) |
| 54 | } |
| 55 | }) |
| 56 | .filter(None, LevelFilter::Trace) |
| 57 | .parse_filters(filter) |
| 58 | .parse_env("RUST_LOG") |
| 59 | .init(); |
| 60 | } |
| 61 | |
| 62 | #[cfg(feature = "log")] |
| 63 | pub fn setup_logging(filter: &str) { |
| 64 | setup_logging_with_clock(filter, Instant::now) |
| 65 | } |
| 66 | |
| 67 | pub fn create_options() -> (Options, Vec<&'static str>) { |
| 68 | let mut opts = Options::new(); |
| 69 | opts.optflag("h", "help", "print this help menu"); |
| 70 | (opts, Vec::new()) |
| 71 | } |
| 72 | |
| 73 | pub fn parse_options(options: &Options, free: Vec<&str>) -> Matches { |
| 74 | match options.parse(env::args().skip(1)) { |
| 75 | Err(err) => { |
| 76 | println!("{err}"); |
| 77 | process::exit(1) |
| 78 | } |
| 79 | Ok(matches) => { |
| 80 | if matches.opt_present("h") || matches.free.len() != free.len() { |
| 81 | let brief = format!( |
| 82 | "Usage: {} [OPTION]... {}", |
| 83 | env::args().next().unwrap(), |
| 84 | free.join(" ") |
| 85 | ); |
| 86 | print!("{}", options.usage(&brief)); |
| 87 | process::exit((matches.free.len() != free.len()) as _); |
| 88 | } |
| 89 | matches |
| 90 | } |
| 91 | } |
| 92 | } |
| 93 | |
| 94 | pub fn add_tuntap_options(opts: &mut Options, _free: &mut [&str]) { |
| 95 | opts.optopt("", "tun", "TUN interface to use", "tun0"); |
| 96 | opts.optopt("", "tap", "TAP interface to use", "tap0"); |
| 97 | } |
| 98 | |
| 99 | #[cfg(feature = "phy-tuntap_interface")] |
| 100 | pub fn parse_tuntap_options(matches: &mut Matches) -> TunTapInterface { |
| 101 | let tun = matches.opt_str("tun"); |
| 102 | let tap = matches.opt_str("tap"); |
| 103 | match (tun, tap) { |
| 104 | (Some(tun), None) => TunTapInterface::new(&tun, Medium::Ip).unwrap(), |
| 105 | (None, Some(tap)) => TunTapInterface::new(&tap, Medium::Ethernet).unwrap(), |
| 106 | _ => panic!("You must specify exactly one of --tun or --tap"), |
| 107 | } |
| 108 | } |
| 109 | |
| 110 | pub fn add_middleware_options(opts: &mut Options, _free: &mut [&str]) { |
| 111 | opts.optopt("", "pcap", "Write a packet capture file", "FILE"); |
| 112 | opts.optopt( |
| 113 | "", |
| 114 | "drop-chance", |
| 115 | "Chance of dropping a packet (%)", |
| 116 | "CHANCE", |
| 117 | ); |
| 118 | opts.optopt( |
| 119 | "", |
| 120 | "corrupt-chance", |
| 121 | "Chance of corrupting a packet (%)", |
| 122 | "CHANCE", |
| 123 | ); |
| 124 | opts.optopt( |
| 125 | "", |
| 126 | "size-limit", |
| 127 | "Drop packets larger than given size (octets)", |
| 128 | "SIZE", |
| 129 | ); |
| 130 | opts.optopt( |
| 131 | "", |
| 132 | "tx-rate-limit", |
| 133 | "Drop packets after transmit rate exceeds given limit \ |
| 134 | (packets per interval)", |
| 135 | "RATE", |
| 136 | ); |
| 137 | opts.optopt( |
| 138 | "", |
| 139 | "rx-rate-limit", |
| 140 | "Drop packets after transmit rate exceeds given limit \ |
| 141 | (packets per interval)", |
| 142 | "RATE", |
| 143 | ); |
| 144 | opts.optopt( |
| 145 | "", |
| 146 | "shaping-interval", |
| 147 | "Sets the interval for rate limiting (ms)", |
| 148 | "RATE", |
| 149 | ); |
| 150 | } |
| 151 | |
| 152 | pub fn parse_middleware_options<D>( |
| 153 | matches: &mut Matches, |
| 154 | device: D, |
| 155 | loopback: bool, |
| 156 | ) -> FaultInjector<Tracer<PcapWriter<D, Box<dyn io::Write>>>> |
| 157 | where |
| 158 | D: Device, |
| 159 | { |
| 160 | let drop_chance = matches |
| 161 | .opt_str("drop-chance") |
| 162 | .map(|s| u8::from_str(&s).unwrap()) |
| 163 | .unwrap_or(0); |
| 164 | let corrupt_chance = matches |
| 165 | .opt_str("corrupt-chance") |
| 166 | .map(|s| u8::from_str(&s).unwrap()) |
| 167 | .unwrap_or(0); |
| 168 | let size_limit = matches |
| 169 | .opt_str("size-limit") |
| 170 | .map(|s| usize::from_str(&s).unwrap()) |
| 171 | .unwrap_or(0); |
| 172 | let tx_rate_limit = matches |
| 173 | .opt_str("tx-rate-limit") |
| 174 | .map(|s| u64::from_str(&s).unwrap()) |
| 175 | .unwrap_or(0); |
| 176 | let rx_rate_limit = matches |
| 177 | .opt_str("rx-rate-limit") |
| 178 | .map(|s| u64::from_str(&s).unwrap()) |
| 179 | .unwrap_or(0); |
| 180 | let shaping_interval = matches |
| 181 | .opt_str("shaping-interval") |
| 182 | .map(|s| u64::from_str(&s).unwrap()) |
| 183 | .unwrap_or(0); |
| 184 | |
| 185 | let pcap_writer: Box<dyn io::Write> = match matches.opt_str("pcap") { |
| 186 | Some(pcap_filename) => Box::new(File::create(pcap_filename).expect("cannot open file")), |
| 187 | None => Box::new(io::sink()), |
| 188 | }; |
| 189 | |
| 190 | let seed = SystemTime::now() |
| 191 | .duration_since(UNIX_EPOCH) |
| 192 | .unwrap() |
| 193 | .subsec_nanos(); |
| 194 | |
| 195 | let device = PcapWriter::new( |
| 196 | device, |
| 197 | pcap_writer, |
| 198 | if loopback { |
| 199 | PcapMode::TxOnly |
| 200 | } else { |
| 201 | PcapMode::Both |
| 202 | }, |
| 203 | ); |
| 204 | |
| 205 | let device = Tracer::new(device, |_timestamp, _printer| { |
| 206 | #[cfg(feature = "log")] |
| 207 | trace!("{}", _printer); |
| 208 | }); |
| 209 | |
| 210 | let mut device = FaultInjector::new(device, seed); |
| 211 | device.set_drop_chance(drop_chance); |
| 212 | device.set_corrupt_chance(corrupt_chance); |
| 213 | device.set_max_packet_size(size_limit); |
| 214 | device.set_max_tx_rate(tx_rate_limit); |
| 215 | device.set_max_rx_rate(rx_rate_limit); |
| 216 | device.set_bucket_interval(Duration::from_millis(shaping_interval)); |
| 217 | device |
| 218 | } |