浏览代码

Add --pcap option to all our examples.

Also, generally reorganize and clean up option handling.
whitequark 7 年之前
父节点
当前提交
c799bfc398
共有 5 个文件被更改,包括 133 次插入40 次删除
  1. 13 4
      examples/client.rs
  2. 12 1
      examples/loopback.rs
  3. 25 11
      examples/ping.rs
  4. 9 2
      examples/server.rs
  5. 74 22
      examples/utils.rs

+ 13 - 4
examples/client.rs

@@ -15,10 +15,19 @@ use smoltcp::socket::{AsSocket, SocketSet};
 use smoltcp::socket::{TcpSocket, TcpSocketBuffer};
 
 fn main() {
-    utils::setup_logging();
-    let (device, args) = utils::setup_device(&["ADDRESS", "PORT"]);
-    let address = IpAddress::from_str(&args[0]).expect("invalid address format");
-    let port = u16::from_str(&args[1]).expect("invalid port format");
+    utils::setup_logging("");
+
+    let (mut opts, mut free) = utils::create_options();
+    utils::add_tap_options(&mut opts, &mut free);
+    utils::add_middleware_options(&mut opts, &mut free);
+    free.push("ADDRESS");
+    free.push("PORT");
+
+    let mut matches = utils::parse_options(&opts, free);
+    let device = utils::parse_tap_options(&mut matches);
+    let device = utils::parse_middleware_options(&mut matches, device, /*loopback=*/false);
+    let address = IpAddress::from_str(&matches.free[0]).expect("invalid address format");
+    let port = u16::from_str(&matches.free[1]).expect("invalid port format");
 
     let startup_time = Instant::now();
 

+ 12 - 1
examples/loopback.rs

@@ -73,12 +73,23 @@ fn main() {
     #[cfg(feature = "std")]
     {
         let clock = clock.clone();
-        utils::setup_logging_with_clock(move || clock.elapsed());
+        utils::setup_logging_with_clock("", move || clock.elapsed());
     }
 
     let mut device = Loopback::new();
     let mut device = EthernetTracer::new(device, |_timestamp, printer| trace!("{}", printer));
 
+    #[cfg(feature = "std")]
+    let mut device = {
+        let (mut opts, mut free) = utils::create_options();
+        utils::add_middleware_options(&mut opts, &mut free);
+
+        let mut matches = utils::parse_options(&opts, free);
+        let device = utils::parse_middleware_options(&mut matches, device, /*loopback=*/true);
+
+        device
+    };
+
     let mut arp_cache_entries: [_; 8] = Default::default();
     let mut arp_cache = SliceArpCache::new(&mut arp_cache_entries[..]);
 

+ 25 - 11
examples/ping.rs

@@ -19,19 +19,33 @@ use smoltcp::socket::{RawSocket, RawSocketBuffer, RawPacketBuffer};
 use std::collections::HashMap;
 use byteorder::{ByteOrder, NetworkEndian};
 
-const PING_INTERVAL_S: u64 = 1;
-const PING_TIMEOUT_S: u64 = 5;
-const PINGS_TO_SEND: usize = 4;
-
 fn main() {
-    utils::setup_logging();
-    let (device, args) = utils::setup_device(&["ADDRESS"]);
+    utils::setup_logging("warn");
+
+    let (mut opts, mut free) = utils::create_options();
+    utils::add_tap_options(&mut opts, &mut free);
+    utils::add_middleware_options(&mut opts, &mut free);
+    opts.optopt("c", "count", "Amount of echo request packets to send (default: 4)", "COUNT");
+    opts.optopt("i", "interval",
+                "Interval between successive packets sent (seconds) (default: 1)", "INTERVAL");
+    opts.optopt("", "timeout",
+                "Maximum wait duration for an echo response packet (seconds) (default: 5)",
+                "TIMEOUT");
+    free.push("ADDRESS");
+
+    let mut matches = utils::parse_options(&opts, free);
+    let device = utils::parse_tap_options(&mut matches);
+    let device = utils::parse_middleware_options(&mut matches, device, /*loopback=*/false);
+    let address  = Ipv4Address::from_str(&matches.free[0]).expect("invalid address format");
+    let count    = matches.opt_str("count").map(|s| usize::from_str(&s).unwrap()).unwrap_or(4);
+    let interval = matches.opt_str("interval").map(|s| u64::from_str(&s).unwrap()).unwrap_or(1);
+    let timeout  = matches.opt_str("timeout").map(|s| u64::from_str(&s).unwrap()).unwrap_or(5);
 
     let startup_time = Instant::now();
 
     let arp_cache = SliceArpCache::new(vec![Default::default(); 8]);
 
-    let remote_addr = Ipv4Address::from_str(&args[0]).unwrap();
+    let remote_addr = address;
     let local_addr  = Ipv4Address::new(192, 168, 69, 1);
 
     let raw_rx_buffer = RawSocketBuffer::new(vec![RawPacketBuffer::new(vec![0; 256])]);
@@ -61,11 +75,11 @@ fn main() {
             let timestamp_us = (timestamp.as_secs() * 1000000) +
                 (timestamp.subsec_nanos() / 1000) as u64;
 
-            if seq_no == PINGS_TO_SEND as u16 && waiting_queue.is_empty() {
+            if seq_no == count as u16 && waiting_queue.is_empty() {
                 break;
             }
 
-            if socket.can_send() && seq_no < PINGS_TO_SEND as u16 && send_next <= timestamp {
+            if socket.can_send() && seq_no < count as u16 && send_next <= timestamp {
                 NetworkEndian::write_u64(&mut echo_payload, timestamp_us);
                 let icmp_repr = Icmpv4Repr::EchoRequest {
                     ident: 1,
@@ -91,7 +105,7 @@ fn main() {
 
                 waiting_queue.insert(seq_no, timestamp);
                 seq_no += 1;
-                send_next += Duration::new(PING_INTERVAL_S, 0);
+                send_next += Duration::new(interval, 0);
             }
 
             if socket.can_recv() {
@@ -117,7 +131,7 @@ fn main() {
             }
 
             waiting_queue.retain(|seq, from| {
-                if (timestamp - *from).as_secs() < PING_TIMEOUT_S {
+                if (timestamp - *from).as_secs() < timeout {
                     true
                 } else {
                     println!("From {} icmp_seq={} timeout", remote_addr, seq);

+ 9 - 2
examples/server.rs

@@ -16,8 +16,15 @@ use smoltcp::socket::{UdpSocket, UdpSocketBuffer, UdpPacketBuffer};
 use smoltcp::socket::{TcpSocket, TcpSocketBuffer};
 
 fn main() {
-    utils::setup_logging();
-    let (device, _args) = utils::setup_device(&[]);
+    utils::setup_logging("");
+
+    let (mut opts, mut free) = utils::create_options();
+    utils::add_tap_options(&mut opts, &mut free);
+    utils::add_middleware_options(&mut opts, &mut free);
+
+    let mut matches = utils::parse_options(&opts, free);
+    let device = utils::parse_tap_options(&mut matches);
+    let device = utils::parse_middleware_options(&mut matches, device, /*loopback=*/false);
 
     let startup_time = Instant::now();
 

+ 74 - 22
examples/utils.rs

@@ -1,14 +1,19 @@
+use std::cell::RefCell;
 use std::str::{self, FromStr};
-use std::env;
+use std::rc::Rc;
+use std::io;
+use std::fs::File;
 use std::time::{Instant, Duration, SystemTime, UNIX_EPOCH};
+use std::env;
 use std::process;
 use log::{LogLevel, LogLevelFilter, LogRecord};
 use env_logger::LogBuilder;
-use getopts;
+use getopts::{Options, Matches};
 
-use smoltcp::phy::{EthernetTracer, FaultInjector, TapInterface};
+use smoltcp::phy::{Device, EthernetTracer, FaultInjector, TapInterface};
+use smoltcp::phy::{PcapWriter, PcapSink, PcapMode, PcapLinkType};
 
-pub fn setup_logging_with_clock<F>(since_startup: F)
+pub fn setup_logging_with_clock<F>(filter: &str, since_startup: F)
         where F: Fn() -> u64 + Send + Sync + 'static {
     LogBuilder::new()
         .format(move |record: &LogRecord| {
@@ -28,22 +33,67 @@ pub fn setup_logging_with_clock<F>(since_startup: F)
             }
         })
         .filter(None, LogLevelFilter::Trace)
+        .parse(filter)
+        .parse(&env::var("RUST_LOG").unwrap_or("".to_owned()))
         .init()
         .unwrap();
 }
 
-pub fn setup_logging() {
+pub fn setup_logging(filter: &str) {
     let startup_at = Instant::now();
-    setup_logging_with_clock(move  || {
+    setup_logging_with_clock(filter, move  || {
         let elapsed = Instant::now().duration_since(startup_at);
         elapsed.as_secs() * 1000 + (elapsed.subsec_nanos() / 1000000) as u64
     })
 }
 
-pub fn setup_device(more_args: &[&str])
-        -> (FaultInjector<EthernetTracer<TapInterface>>,
-            Vec<String>) {
-    let mut opts = getopts::Options::new();
+struct Dispose;
+
+impl io::Write for Dispose {
+    fn write(&mut self, data: &[u8]) -> io::Result<usize> {
+        Ok(data.len())
+    }
+
+    fn flush(&mut self) -> io::Result<()> {
+        Ok(())
+    }
+}
+
+pub fn create_options() -> (Options, Vec<&'static str>) {
+    let mut opts = Options::new();
+    opts.optflag("h", "help", "print this help menu");
+    (opts, Vec::new())
+}
+
+pub fn parse_options(options: &Options, free: Vec<&str>) -> Matches {
+    match options.parse(env::args().skip(1)) {
+        Err(err) => {
+            println!("{}", err);
+            process::exit(1)
+        }
+        Ok(matches) => {
+            if matches.opt_present("h") || matches.free.len() != free.len() {
+                let brief = format!("Usage: {} [OPTION]... {}",
+                                    env::args().nth(0).unwrap(), free.join(" "));
+                print!("{}", options.usage(&brief));
+                process::exit(if matches.free.len() != free.len() { 1 } else { 0 })
+            }
+            matches
+        }
+    }
+}
+
+pub fn add_tap_options(_opts: &mut Options, free: &mut Vec<&str>) {
+    free.push("INTERFACE");
+}
+
+pub fn parse_tap_options(matches: &mut Matches) -> TapInterface {
+    let interface = matches.free.remove(0);
+    TapInterface::new(&interface).unwrap()
+}
+
+pub fn add_middleware_options(opts: &mut Options, _free: &mut Vec<&str>) {
+    opts.optopt("", "pcap", "Write a packet capture file", "FILE");
     opts.optopt("", "drop-chance", "Chance of dropping a packet (%)", "CHANCE");
     opts.optopt("", "corrupt-chance", "Chance of corrupting a packet (%)", "CHANCE");
     opts.optopt("", "size-limit", "Drop packets larger than given size (octets)", "SIZE");
@@ -52,16 +102,10 @@ pub fn setup_device(more_args: &[&str])
     opts.optopt("", "rx-rate-limit", "Drop packets after transmit rate exceeds given limit \
                                       (packets per interval)", "RATE");
     opts.optopt("", "shaping-interval", "Sets the interval for rate limiting (ms)", "RATE");
-    opts.optflag("h", "help", "print this help menu");
+}
 
-    let matches = opts.parse(env::args().skip(1)).unwrap();
-    if matches.opt_present("h") || matches.free.len() != more_args.len() + 1 {
-        let brief = format!("Usage: {} INTERFACE {} [options]",
-                            env::args().nth(0).unwrap(),
-                            more_args.join(" "));
-        print!("{}", opts.usage(&brief));
-        process::exit(if matches.free.len() != more_args.len() + 1 { 1 } else { 0 });
-    }
+pub fn parse_middleware_options<D: Device>(matches: &mut Matches, device: D, loopback: bool)
+        -> FaultInjector<EthernetTracer<PcapWriter<D, Rc<PcapSink>>>> {
     let drop_chance    = u8::from_str(&matches.opt_str("drop-chance")
                                              .unwrap_or("0".to_string())).unwrap();
     let corrupt_chance = u8::from_str(&matches.opt_str("corrupt-chance")
@@ -75,9 +119,18 @@ pub fn setup_device(more_args: &[&str])
     let shaping_interval = u32::from_str(&matches.opt_str("shaping-interval")
                                                  .unwrap_or("0".to_string())).unwrap();
 
+    let pcap_writer: Box<io::Write>;
+    if let Some(pcap_filename) = matches.opt_str("pcap") {
+        pcap_writer = Box::new(File::create(pcap_filename).expect("cannot open file"))
+    } else {
+        pcap_writer = Box::new(Dispose)
+    }
+
     let seed = SystemTime::now().duration_since(UNIX_EPOCH).unwrap().subsec_nanos();
 
-    let device = TapInterface::new(&matches.free[0]).unwrap();
+    let device = PcapWriter::new(device, Rc::new(RefCell::new(pcap_writer)) as Rc<PcapSink>,
+                                 if loopback { PcapMode::TxOnly } else { PcapMode::Both },
+                                 PcapLinkType::Ethernet);
     let device = EthernetTracer::new(device, |_timestamp, printer| trace!("{}", printer));
     let mut device = FaultInjector::new(device, seed);
     device.set_drop_chance(drop_chance);
@@ -86,6 +139,5 @@ pub fn setup_device(more_args: &[&str])
     device.set_max_tx_rate(tx_rate_limit);
     device.set_max_rx_rate(rx_rate_limit);
     device.set_bucket_interval(Duration::from_millis(shaping_interval as u64));
-
-    (device, matches.free[1..].to_owned())
+    device
 }