utils.rs 6.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149
  1. #![allow(dead_code)]
  2. use std::cell::RefCell;
  3. use std::str::{self, FromStr};
  4. use std::rc::Rc;
  5. use std::io;
  6. use std::fs::File;
  7. use std::time::{SystemTime, UNIX_EPOCH};
  8. use std::env;
  9. use std::process;
  10. #[cfg(feature = "log")]
  11. use log::{LogLevel, LogLevelFilter, LogRecord};
  12. #[cfg(feature = "log")]
  13. use env_logger::LogBuilder;
  14. use getopts::{Options, Matches};
  15. use smoltcp::phy::{Device, EthernetTracer, FaultInjector};
  16. #[cfg(feature = "phy-tap_interface")]
  17. use smoltcp::phy::TapInterface;
  18. use smoltcp::phy::{PcapWriter, PcapSink, PcapMode, PcapLinkType};
  19. use smoltcp::phy::RawSocket;
  20. use smoltcp::time::{Duration, Instant};
  21. #[cfg(feature = "log")]
  22. pub fn setup_logging_with_clock<F>(filter: &str, since_startup: F)
  23. where F: Fn() -> Instant + Send + Sync + 'static {
  24. LogBuilder::new()
  25. .format(move |record: &LogRecord| {
  26. let elapsed = since_startup();
  27. let timestamp = format!("[{}]", elapsed);
  28. if record.target().starts_with("smoltcp::") {
  29. format!("\x1b[0m{} ({}): {}\x1b[0m", timestamp,
  30. record.target().replace("smoltcp::", ""), record.args())
  31. } else if record.level() == LogLevel::Trace {
  32. let message = format!("{}", record.args());
  33. format!("\x1b[37m{} {}\x1b[0m", timestamp,
  34. message.replace("\n", "\n "))
  35. } else {
  36. format!("\x1b[32m{} ({}): {}\x1b[0m", timestamp,
  37. record.target(), record.args())
  38. }
  39. })
  40. .filter(None, LogLevelFilter::Trace)
  41. .parse(filter)
  42. .parse(&env::var("RUST_LOG").unwrap_or("".to_owned()))
  43. .init()
  44. .unwrap();
  45. }
  46. #[cfg(feature = "log")]
  47. pub fn setup_logging(filter: &str) {
  48. setup_logging_with_clock(filter, move || {
  49. Instant::now()
  50. })
  51. }
  52. pub fn create_options() -> (Options, Vec<&'static str>) {
  53. let mut opts = Options::new();
  54. opts.optflag("h", "help", "print this help menu");
  55. (opts, Vec::new())
  56. }
  57. pub fn parse_options(options: &Options, free: Vec<&str>) -> Matches {
  58. match options.parse(env::args().skip(1)) {
  59. Err(err) => {
  60. println!("{}", err);
  61. process::exit(1)
  62. }
  63. Ok(matches) => {
  64. if matches.opt_present("h") || matches.free.len() != free.len() {
  65. let brief = format!("Usage: {} [OPTION]... {}",
  66. env::args().nth(0).unwrap(), free.join(" "));
  67. print!("{}", options.usage(&brief));
  68. process::exit(if matches.free.len() != free.len() { 1 } else { 0 })
  69. }
  70. matches
  71. }
  72. }
  73. }
  74. pub fn add_tap_options(_opts: &mut Options, free: &mut Vec<&str>) {
  75. free.push("INTERFACE");
  76. }
  77. #[cfg(feature = "phy-tap_interface")]
  78. pub fn parse_tap_options(matches: &mut Matches) -> TapInterface {
  79. let interface = matches.free.remove(0);
  80. TapInterface::new(&interface).unwrap()
  81. }
  82. pub fn parse_raw_socket_options(matches: &mut Matches) -> RawSocket {
  83. let interface = matches.free.remove(0);
  84. RawSocket::new(&interface).unwrap()
  85. }
  86. pub fn add_middleware_options(opts: &mut Options, _free: &mut Vec<&str>) {
  87. opts.optopt("", "pcap", "Write a packet capture file", "FILE");
  88. opts.optopt("", "drop-chance", "Chance of dropping a packet (%)", "CHANCE");
  89. opts.optopt("", "corrupt-chance", "Chance of corrupting a packet (%)", "CHANCE");
  90. opts.optopt("", "size-limit", "Drop packets larger than given size (octets)", "SIZE");
  91. opts.optopt("", "tx-rate-limit", "Drop packets after transmit rate exceeds given limit \
  92. (packets per interval)", "RATE");
  93. opts.optopt("", "rx-rate-limit", "Drop packets after transmit rate exceeds given limit \
  94. (packets per interval)", "RATE");
  95. opts.optopt("", "shaping-interval", "Sets the interval for rate limiting (ms)", "RATE");
  96. }
  97. pub fn parse_middleware_options<D>(matches: &mut Matches, device: D, loopback: bool)
  98. -> FaultInjector<EthernetTracer<PcapWriter<D, Rc<PcapSink>>>>
  99. where D: for<'a> Device<'a>
  100. {
  101. let drop_chance = matches.opt_str("drop-chance").map(|s| u8::from_str(&s).unwrap())
  102. .unwrap_or(0);
  103. let corrupt_chance = matches.opt_str("corrupt-chance").map(|s| u8::from_str(&s).unwrap())
  104. .unwrap_or(0);
  105. let size_limit = matches.opt_str("size-limit").map(|s| usize::from_str(&s).unwrap())
  106. .unwrap_or(0);
  107. let tx_rate_limit = matches.opt_str("tx-rate-limit").map(|s| u64::from_str(&s).unwrap())
  108. .unwrap_or(0);
  109. let rx_rate_limit = matches.opt_str("rx-rate-limit").map(|s| u64::from_str(&s).unwrap())
  110. .unwrap_or(0);
  111. let shaping_interval = matches.opt_str("shaping-interval").map(|s| u64::from_str(&s).unwrap())
  112. .unwrap_or(0);
  113. let pcap_writer: Box<io::Write>;
  114. if let Some(pcap_filename) = matches.opt_str("pcap") {
  115. pcap_writer = Box::new(File::create(pcap_filename).expect("cannot open file"))
  116. } else {
  117. pcap_writer = Box::new(io::sink())
  118. }
  119. let seed = SystemTime::now().duration_since(UNIX_EPOCH).unwrap().subsec_nanos();
  120. let device = PcapWriter::new(device, Rc::new(RefCell::new(pcap_writer)) as Rc<PcapSink>,
  121. if loopback { PcapMode::TxOnly } else { PcapMode::Both },
  122. PcapLinkType::Ethernet);
  123. let device = EthernetTracer::new(device, |_timestamp, _printer| {
  124. #[cfg(feature = "log")]
  125. trace!("{}", _printer);
  126. });
  127. let mut device = FaultInjector::new(device, seed);
  128. device.set_drop_chance(drop_chance);
  129. device.set_corrupt_chance(corrupt_chance);
  130. device.set_max_packet_size(size_limit);
  131. device.set_max_tx_rate(tx_rate_limit);
  132. device.set_max_rx_rate(rx_rate_limit);
  133. device.set_bucket_interval(Duration::from_millis(shaping_interval));
  134. device
  135. }