浏览代码

Update fault injector so that rate limiting works without std.

whitequark 7 年之前
父节点
当前提交
d73c2bdb81
共有 2 个文件被更改,包括 29 次插入62 次删除
  1. 14 14
      examples/utils.rs
  2. 15 48
      src/phy/fault_injector.rs

+ 14 - 14
examples/utils.rs

@@ -3,7 +3,7 @@ use std::str::{self, FromStr};
 use std::rc::Rc;
 use std::rc::Rc;
 use std::io;
 use std::io;
 use std::fs::File;
 use std::fs::File;
-use std::time::{Instant, Duration, SystemTime, UNIX_EPOCH};
+use std::time::{Instant, SystemTime, UNIX_EPOCH};
 use std::env;
 use std::env;
 use std::process;
 use std::process;
 use log::{LogLevel, LogLevelFilter, LogRecord};
 use log::{LogLevel, LogLevelFilter, LogRecord};
@@ -106,18 +106,18 @@ pub fn add_middleware_options(opts: &mut Options, _free: &mut Vec<&str>) {
 
 
 pub fn parse_middleware_options<D: Device>(matches: &mut Matches, device: D, loopback: bool)
 pub fn parse_middleware_options<D: Device>(matches: &mut Matches, device: D, loopback: bool)
         -> FaultInjector<EthernetTracer<PcapWriter<D, Rc<PcapSink>>>> {
         -> 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")
-                                             .unwrap_or("0".to_string())).unwrap();
-    let size_limit = usize::from_str(&matches.opt_str("size-limit")
-                                             .unwrap_or("0".to_string())).unwrap();
-    let tx_rate_limit = u64::from_str(&matches.opt_str("tx-rate-limit")
-                                              .unwrap_or("0".to_string())).unwrap();
-    let rx_rate_limit = u64::from_str(&matches.opt_str("rx-rate-limit")
-                                              .unwrap_or("0".to_string())).unwrap();
-    let shaping_interval = u32::from_str(&matches.opt_str("shaping-interval")
-                                                 .unwrap_or("0".to_string())).unwrap();
+    let drop_chance      = matches.opt_str("drop-chance").map(|s| u8::from_str(&s).unwrap())
+                                  .unwrap_or(0);
+    let corrupt_chance   = matches.opt_str("corrupt-chance").map(|s| u8::from_str(&s).unwrap())
+                                  .unwrap_or(0);
+    let size_limit       = matches.opt_str("size-limit").map(|s| usize::from_str(&s).unwrap())
+                                  .unwrap_or(0);
+    let tx_rate_limit    = matches.opt_str("tx-rate-limit").map(|s| u64::from_str(&s).unwrap())
+                                  .unwrap_or(0);
+    let rx_rate_limit    = matches.opt_str("rx-rate-limit").map(|s| u64::from_str(&s).unwrap())
+                                  .unwrap_or(0);
+    let shaping_interval = matches.opt_str("shaping-interval").map(|s| u64::from_str(&s).unwrap())
+                                  .unwrap_or(0);
 
 
     let pcap_writer: Box<io::Write>;
     let pcap_writer: Box<io::Write>;
     if let Some(pcap_filename) = matches.opt_str("pcap") {
     if let Some(pcap_filename) = matches.opt_str("pcap") {
@@ -138,6 +138,6 @@ pub fn parse_middleware_options<D: Device>(matches: &mut Matches, device: D, loo
     device.set_max_packet_size(size_limit);
     device.set_max_packet_size(size_limit);
     device.set_max_tx_rate(tx_rate_limit);
     device.set_max_tx_rate(tx_rate_limit);
     device.set_max_rx_rate(rx_rate_limit);
     device.set_max_rx_rate(rx_rate_limit);
-    device.set_bucket_interval(Duration::from_millis(shaping_interval as u64));
+    device.set_bucket_interval(shaping_interval);
     device
     device
 }
 }

+ 15 - 48
src/phy/fault_injector.rs

@@ -1,6 +1,3 @@
-#[cfg(feature = "std")]
-use std::time::{Instant, Duration};
-
 use Error;
 use Error;
 use super::{DeviceLimits, Device};
 use super::{DeviceLimits, Device};
 
 
@@ -24,22 +21,16 @@ struct Config {
     drop_pct:    u8,
     drop_pct:    u8,
     reorder_pct: u8,
     reorder_pct: u8,
     max_size:    usize,
     max_size:    usize,
-    #[cfg(feature = "std")]
     max_tx_rate: u64,
     max_tx_rate: u64,
-    #[cfg(feature = "std")]
     max_rx_rate: u64,
     max_rx_rate: u64,
-    #[cfg(feature = "std")]
-    interval:    Duration,
+    interval:    u64,
 }
 }
 
 
 #[derive(Debug, Clone, Copy)]
 #[derive(Debug, Clone, Copy)]
 struct State {
 struct State {
     rng_seed:    u32,
     rng_seed:    u32,
-    #[cfg(feature = "std")]
-    refilled_at: Instant,
-    #[cfg(feature = "std")]
+    refilled_at: u64,
     tx_bucket:   u64,
     tx_bucket:   u64,
-    #[cfg(feature = "std")]
     rx_bucket:   u64,
     rx_bucket:   u64,
 }
 }
 
 
@@ -56,20 +47,18 @@ impl State {
         buffer[index] ^= bit;
         buffer[index] ^= bit;
     }
     }
 
 
-    #[cfg(feature = "std")]
-    fn refill(&mut self, config: &Config) {
-        if self.refilled_at.elapsed() > config.interval {
+    fn refill(&mut self, config: &Config, timestamp: u64) {
+        if self.refilled_at - timestamp > config.interval {
             self.tx_bucket = config.max_tx_rate;
             self.tx_bucket = config.max_tx_rate;
             self.rx_bucket = config.max_rx_rate;
             self.rx_bucket = config.max_rx_rate;
-            self.refilled_at = Instant::now();
+            self.refilled_at = timestamp;
         }
         }
     }
     }
 
 
-    #[cfg(feature = "std")]
-    fn maybe_transmit(&mut self, config: &Config) -> bool {
+    fn maybe_transmit(&mut self, config: &Config, timestamp: u64) -> bool {
         if config.max_tx_rate == 0 { return true }
         if config.max_tx_rate == 0 { return true }
 
 
-        self.refill(config);
+        self.refill(config, timestamp);
         if self.tx_bucket > 0 {
         if self.tx_bucket > 0 {
             self.tx_bucket -= 1;
             self.tx_bucket -= 1;
             true
             true
@@ -78,16 +67,10 @@ impl State {
         }
         }
     }
     }
 
 
-    #[cfg(not(feature = "std"))]
-    fn maybe_transmit(&mut self, _config: &Config) -> bool {
-        true
-    }
-
-    #[cfg(feature = "std")]
-    fn maybe_receive(&mut self, config: &Config) -> bool {
+    fn maybe_receive(&mut self, config: &Config, timestamp: u64) -> bool {
         if config.max_rx_rate == 0 { return true }
         if config.max_rx_rate == 0 { return true }
 
 
-        self.refill(config);
+        self.refill(config, timestamp);
         if self.rx_bucket > 0 {
         if self.rx_bucket > 0 {
             self.rx_bucket -= 1;
             self.rx_bucket -= 1;
             true
             true
@@ -95,11 +78,6 @@ impl State {
             false
             false
         }
         }
     }
     }
-
-    #[cfg(not(feature = "std"))]
-    fn maybe_receive(&mut self, _config: &Config) -> bool {
-        true
-    }
 }
 }
 
 
 /// A fault injector device.
 /// A fault injector device.
@@ -117,17 +95,12 @@ pub struct FaultInjector<D: Device> {
 impl<D: Device> FaultInjector<D> {
 impl<D: Device> FaultInjector<D> {
     /// Create a fault injector device, using the given random number generator seed.
     /// Create a fault injector device, using the given random number generator seed.
     pub fn new(inner: D, seed: u32) -> FaultInjector<D> {
     pub fn new(inner: D, seed: u32) -> FaultInjector<D> {
-        #[cfg(feature = "std")]
         let state = State {
         let state = State {
             rng_seed:    seed,
             rng_seed:    seed,
-            refilled_at: Instant::now(),
+            refilled_at: 0,
             tx_bucket:   0,
             tx_bucket:   0,
             rx_bucket:   0,
             rx_bucket:   0,
         };
         };
-        #[cfg(not(feature = "std"))]
-        let state = State {
-            rng_seed:    seed,
-        };
         FaultInjector {
         FaultInjector {
             inner: inner,
             inner: inner,
             state: state,
             state: state,
@@ -156,20 +129,17 @@ impl<D: Device> FaultInjector<D> {
     }
     }
 
 
     /// Return the maximum packet transmission rate, in packets per second.
     /// Return the maximum packet transmission rate, in packets per second.
-    #[cfg(feature = "std")]
     pub fn max_tx_rate(&self) -> u64 {
     pub fn max_tx_rate(&self) -> u64 {
         self.config.max_rx_rate
         self.config.max_rx_rate
     }
     }
 
 
     /// Return the maximum packet reception rate, in packets per second.
     /// Return the maximum packet reception rate, in packets per second.
-    #[cfg(feature = "std")]
     pub fn max_rx_rate(&self) -> u64 {
     pub fn max_rx_rate(&self) -> u64 {
         self.config.max_tx_rate
         self.config.max_tx_rate
     }
     }
 
 
     /// Return the interval for packet rate limiting, in milliseconds.
     /// Return the interval for packet rate limiting, in milliseconds.
-    #[cfg(feature = "std")]
-    pub fn bucket_interval(&self) -> Duration {
+    pub fn bucket_interval(&self) -> u64 {
         self.config.interval
         self.config.interval
     }
     }
 
 
@@ -197,21 +167,18 @@ impl<D: Device> FaultInjector<D> {
     }
     }
 
 
     /// Set the maximum packet transmission rate, in packets per interval.
     /// Set the maximum packet transmission rate, in packets per interval.
-    #[cfg(feature = "std")]
     pub fn set_max_tx_rate(&mut self, rate: u64) {
     pub fn set_max_tx_rate(&mut self, rate: u64) {
         self.config.max_tx_rate = rate
         self.config.max_tx_rate = rate
     }
     }
 
 
     /// Set the maximum packet reception rate, in packets per interval.
     /// Set the maximum packet reception rate, in packets per interval.
-    #[cfg(feature = "std")]
     pub fn set_max_rx_rate(&mut self, rate: u64) {
     pub fn set_max_rx_rate(&mut self, rate: u64) {
         self.config.max_rx_rate = rate
         self.config.max_rx_rate = rate
     }
     }
 
 
     /// Set the interval for packet rate limiting, in milliseconds.
     /// Set the interval for packet rate limiting, in milliseconds.
-    #[cfg(feature = "std")]
-    pub fn set_bucket_interval(&mut self, interval: Duration) {
-        self.state.refilled_at = Instant::now() - self.config.interval;
+    pub fn set_bucket_interval(&mut self, interval: u64) {
+        self.state.refilled_at = 0;
         self.config.interval = interval
         self.config.interval = interval
     }
     }
 }
 }
@@ -243,7 +210,7 @@ impl<D: Device> Device for FaultInjector<D>
             net_trace!("rx: dropping a packet that is too large");
             net_trace!("rx: dropping a packet that is too large");
             return Err(Error::Exhausted)
             return Err(Error::Exhausted)
         }
         }
-        if !self.state.maybe_receive(&self.config) {
+        if !self.state.maybe_receive(&self.config, timestamp) {
             net_trace!("rx: dropping a packet because of rate limiting");
             net_trace!("rx: dropping a packet because of rate limiting");
             return Err(Error::Exhausted)
             return Err(Error::Exhausted)
         }
         }
@@ -258,7 +225,7 @@ impl<D: Device> Device for FaultInjector<D>
         } else if self.config.max_size > 0 && length > self.config.max_size {
         } else if self.config.max_size > 0 && length > self.config.max_size {
             net_trace!("tx: dropping a packet that is too large");
             net_trace!("tx: dropping a packet that is too large");
             buffer = None;
             buffer = None;
-        } else if !self.state.maybe_transmit(&self.config) {
+        } else if !self.state.maybe_transmit(&self.config, timestamp) {
             net_trace!("tx: dropping a packet because of rate limiting");
             net_trace!("tx: dropping a packet because of rate limiting");
             buffer = None;
             buffer = None;
         } else {
         } else {