Skip to content

Commit d73c2bd

Browse files
whitequarkwhitequark
committedJul 23, 2017
Update fault injector so that rate limiting works without std.
·
v0.12.0v0.4.0
1 parent c799bfc commit d73c2bd

File tree

2 files changed

+29
-62
lines changed

2 files changed

+29
-62
lines changed
 

‎examples/utils.rs

+14-14
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@ use std::str::{self, FromStr};
33
use std::rc::Rc;
44
use std::io;
55
use std::fs::File;
6-
use std::time::{Instant, Duration, SystemTime, UNIX_EPOCH};
6+
use std::time::{Instant, SystemTime, UNIX_EPOCH};
77
use std::env;
88
use std::process;
99
use log::{LogLevel, LogLevelFilter, LogRecord};
@@ -106,18 +106,18 @@ pub fn add_middleware_options(opts: &mut Options, _free: &mut Vec<&str>) {
106106

107107
pub fn parse_middleware_options<D: Device>(matches: &mut Matches, device: D, loopback: bool)
108108
-> FaultInjector<EthernetTracer<PcapWriter<D, Rc<PcapSink>>>> {
109-
let drop_chance = u8::from_str(&matches.opt_str("drop-chance")
110-
.unwrap_or("0".to_string())).unwrap();
111-
let corrupt_chance = u8::from_str(&matches.opt_str("corrupt-chance")
112-
.unwrap_or("0".to_string())).unwrap();
113-
let size_limit = usize::from_str(&matches.opt_str("size-limit")
114-
.unwrap_or("0".to_string())).unwrap();
115-
let tx_rate_limit = u64::from_str(&matches.opt_str("tx-rate-limit")
116-
.unwrap_or("0".to_string())).unwrap();
117-
let rx_rate_limit = u64::from_str(&matches.opt_str("rx-rate-limit")
118-
.unwrap_or("0".to_string())).unwrap();
119-
let shaping_interval = u32::from_str(&matches.opt_str("shaping-interval")
120-
.unwrap_or("0".to_string())).unwrap();
109+
let drop_chance = matches.opt_str("drop-chance").map(|s| u8::from_str(&s).unwrap())
110+
.unwrap_or(0);
111+
let corrupt_chance = matches.opt_str("corrupt-chance").map(|s| u8::from_str(&s).unwrap())
112+
.unwrap_or(0);
113+
let size_limit = matches.opt_str("size-limit").map(|s| usize::from_str(&s).unwrap())
114+
.unwrap_or(0);
115+
let tx_rate_limit = matches.opt_str("tx-rate-limit").map(|s| u64::from_str(&s).unwrap())
116+
.unwrap_or(0);
117+
let rx_rate_limit = matches.opt_str("rx-rate-limit").map(|s| u64::from_str(&s).unwrap())
118+
.unwrap_or(0);
119+
let shaping_interval = matches.opt_str("shaping-interval").map(|s| u64::from_str(&s).unwrap())
120+
.unwrap_or(0);
121121

122122
let pcap_writer: Box<io::Write>;
123123
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
138138
device.set_max_packet_size(size_limit);
139139
device.set_max_tx_rate(tx_rate_limit);
140140
device.set_max_rx_rate(rx_rate_limit);
141-
device.set_bucket_interval(Duration::from_millis(shaping_interval as u64));
141+
device.set_bucket_interval(shaping_interval);
142142
device
143143
}

‎src/phy/fault_injector.rs

+15-48
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,3 @@
1-
#[cfg(feature = "std")]
2-
use std::time::{Instant, Duration};
3-
41
use Error;
52
use super::{DeviceLimits, Device};
63

@@ -24,22 +21,16 @@ struct Config {
2421
drop_pct: u8,
2522
reorder_pct: u8,
2623
max_size: usize,
27-
#[cfg(feature = "std")]
2824
max_tx_rate: u64,
29-
#[cfg(feature = "std")]
3025
max_rx_rate: u64,
31-
#[cfg(feature = "std")]
32-
interval: Duration,
26+
interval: u64,
3327
}
3428

3529
#[derive(Debug, Clone, Copy)]
3630
struct State {
3731
rng_seed: u32,
38-
#[cfg(feature = "std")]
39-
refilled_at: Instant,
40-
#[cfg(feature = "std")]
32+
refilled_at: u64,
4133
tx_bucket: u64,
42-
#[cfg(feature = "std")]
4334
rx_bucket: u64,
4435
}
4536

@@ -56,20 +47,18 @@ impl State {
5647
buffer[index] ^= bit;
5748
}
5849

59-
#[cfg(feature = "std")]
60-
fn refill(&mut self, config: &Config) {
61-
if self.refilled_at.elapsed() > config.interval {
50+
fn refill(&mut self, config: &Config, timestamp: u64) {
51+
if self.refilled_at - timestamp > config.interval {
6252
self.tx_bucket = config.max_tx_rate;
6353
self.rx_bucket = config.max_rx_rate;
64-
self.refilled_at = Instant::now();
54+
self.refilled_at = timestamp;
6555
}
6656
}
6757

68-
#[cfg(feature = "std")]
69-
fn maybe_transmit(&mut self, config: &Config) -> bool {
58+
fn maybe_transmit(&mut self, config: &Config, timestamp: u64) -> bool {
7059
if config.max_tx_rate == 0 { return true }
7160

72-
self.refill(config);
61+
self.refill(config, timestamp);
7362
if self.tx_bucket > 0 {
7463
self.tx_bucket -= 1;
7564
true
@@ -78,28 +67,17 @@ impl State {
7867
}
7968
}
8069

81-
#[cfg(not(feature = "std"))]
82-
fn maybe_transmit(&mut self, _config: &Config) -> bool {
83-
true
84-
}
85-
86-
#[cfg(feature = "std")]
87-
fn maybe_receive(&mut self, config: &Config) -> bool {
70+
fn maybe_receive(&mut self, config: &Config, timestamp: u64) -> bool {
8871
if config.max_rx_rate == 0 { return true }
8972

90-
self.refill(config);
73+
self.refill(config, timestamp);
9174
if self.rx_bucket > 0 {
9275
self.rx_bucket -= 1;
9376
true
9477
} else {
9578
false
9679
}
9780
}
98-
99-
#[cfg(not(feature = "std"))]
100-
fn maybe_receive(&mut self, _config: &Config) -> bool {
101-
true
102-
}
10381
}
10482

10583
/// A fault injector device.
@@ -117,17 +95,12 @@ pub struct FaultInjector<D: Device> {
11795
impl<D: Device> FaultInjector<D> {
11896
/// Create a fault injector device, using the given random number generator seed.
11997
pub fn new(inner: D, seed: u32) -> FaultInjector<D> {
120-
#[cfg(feature = "std")]
12198
let state = State {
12299
rng_seed: seed,
123-
refilled_at: Instant::now(),
100+
refilled_at: 0,
124101
tx_bucket: 0,
125102
rx_bucket: 0,
126103
};
127-
#[cfg(not(feature = "std"))]
128-
let state = State {
129-
rng_seed: seed,
130-
};
131104
FaultInjector {
132105
inner: inner,
133106
state: state,
@@ -156,20 +129,17 @@ impl<D: Device> FaultInjector<D> {
156129
}
157130

158131
/// Return the maximum packet transmission rate, in packets per second.
159-
#[cfg(feature = "std")]
160132
pub fn max_tx_rate(&self) -> u64 {
161133
self.config.max_rx_rate
162134
}
163135

164136
/// Return the maximum packet reception rate, in packets per second.
165-
#[cfg(feature = "std")]
166137
pub fn max_rx_rate(&self) -> u64 {
167138
self.config.max_tx_rate
168139
}
169140

170141
/// Return the interval for packet rate limiting, in milliseconds.
171-
#[cfg(feature = "std")]
172-
pub fn bucket_interval(&self) -> Duration {
142+
pub fn bucket_interval(&self) -> u64 {
173143
self.config.interval
174144
}
175145

@@ -197,21 +167,18 @@ impl<D: Device> FaultInjector<D> {
197167
}
198168

199169
/// Set the maximum packet transmission rate, in packets per interval.
200-
#[cfg(feature = "std")]
201170
pub fn set_max_tx_rate(&mut self, rate: u64) {
202171
self.config.max_tx_rate = rate
203172
}
204173

205174
/// Set the maximum packet reception rate, in packets per interval.
206-
#[cfg(feature = "std")]
207175
pub fn set_max_rx_rate(&mut self, rate: u64) {
208176
self.config.max_rx_rate = rate
209177
}
210178

211179
/// Set the interval for packet rate limiting, in milliseconds.
212-
#[cfg(feature = "std")]
213-
pub fn set_bucket_interval(&mut self, interval: Duration) {
214-
self.state.refilled_at = Instant::now() - self.config.interval;
180+
pub fn set_bucket_interval(&mut self, interval: u64) {
181+
self.state.refilled_at = 0;
215182
self.config.interval = interval
216183
}
217184
}
@@ -243,7 +210,7 @@ impl<D: Device> Device for FaultInjector<D>
243210
net_trace!("rx: dropping a packet that is too large");
244211
return Err(Error::Exhausted)
245212
}
246-
if !self.state.maybe_receive(&self.config) {
213+
if !self.state.maybe_receive(&self.config, timestamp) {
247214
net_trace!("rx: dropping a packet because of rate limiting");
248215
return Err(Error::Exhausted)
249216
}
@@ -258,7 +225,7 @@ impl<D: Device> Device for FaultInjector<D>
258225
} else if self.config.max_size > 0 && length > self.config.max_size {
259226
net_trace!("tx: dropping a packet that is too large");
260227
buffer = None;
261-
} else if !self.state.maybe_transmit(&self.config) {
228+
} else if !self.state.maybe_transmit(&self.config, timestamp) {
262229
net_trace!("tx: dropping a packet because of rate limiting");
263230
buffer = None;
264231
} else {

0 commit comments

Comments
 (0)
Please sign in to comment.