utils.rs 6.5 KB

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