redis_server_wrapper/
fault_proxy.rs

1//! Byte-level TCP fault-injection proxy for testing Redis client resilience.
2//!
3//! [`FaultProxy`] sits between a test client and an upstream Redis node,
4//! forwarding bytes over TCP while allowing tests to inject faults: per-
5//! direction delay, mid-frame connection drops, chunked writes, and a
6//! black-hole mode. Unlike [`crate::chaos`], which operates on the server
7//! process, this module operates purely on the wire and needs no Docker or
8//! root privileges.
9//!
10//! # Example
11//!
12//! ```no_run
13//! use redis_server_wrapper::{Direction, FaultProxy, RedisServer};
14//! use std::time::Duration;
15//!
16//! # async fn example() {
17//! let server = RedisServer::new().port(6400).start().await.unwrap();
18//! let proxy = FaultProxy::spawn(server.addr()).await.unwrap();
19//!
20//! // Route a client through `proxy.addr()` instead of `server.addr()`.
21//!
22//! // Drop the connection after 8 bytes of the server's response.
23//! proxy.close_after(Direction::UpstreamToClient, 8);
24//!
25//! // ... assert the client sees a clean mid-frame connection error ...
26//!
27//! // Back to clean passthrough for the next connection.
28//! proxy.reset();
29//! # }
30//! ```
31
32use std::net::SocketAddr;
33use std::sync::atomic::{AtomicU64, Ordering};
34use std::sync::{Arc, RwLock};
35use std::time::{Duration, SystemTime, UNIX_EPOCH};
36
37use tokio::io::{AsyncReadExt, AsyncWriteExt};
38use tokio::net::tcp::{OwnedReadHalf, OwnedWriteHalf};
39use tokio::net::{TcpListener, TcpStream, ToSocketAddrs};
40use tokio::task::JoinHandle;
41
42use crate::error::Result;
43
44/// A direction of byte flow through a [`FaultProxy`].
45#[derive(Clone, Copy, Debug, PartialEq, Eq)]
46pub enum Direction {
47    /// Bytes flowing from the test client to the upstream Redis node.
48    ClientToUpstream,
49    /// Bytes flowing from the upstream Redis node to the test client.
50    UpstreamToClient,
51}
52
53impl Direction {
54    fn idx(self) -> usize {
55        match self {
56            Direction::ClientToUpstream => 0,
57            Direction::UpstreamToClient => 1,
58        }
59    }
60}
61
62/// A delay applied before forwarding data in one direction.
63#[derive(Clone, Copy, Debug)]
64pub enum Delay {
65    /// Always sleep for exactly this long before forwarding.
66    Fixed(Duration),
67    /// Sleep for a pseudo-random duration in `[min, max)` before forwarding.
68    ///
69    /// This uses a small local PRNG for jitter, not a cryptographic source --
70    /// it's meant to simulate jittery network latency in tests, not to be
71    /// unpredictable.
72    Random {
73        /// Minimum delay, inclusive.
74        min: Duration,
75        /// Maximum delay, exclusive.
76        max: Duration,
77    },
78}
79
80impl Delay {
81    fn resolve(self) -> Duration {
82        match self {
83            Delay::Fixed(d) => d,
84            Delay::Random { min, max } => {
85                if max <= min {
86                    return min;
87                }
88                let span = (max - min).as_nanos().max(1) as u64;
89                min + Duration::from_nanos(next_pseudo_random() % span)
90            }
91        }
92    }
93}
94
95/// A small local PRNG (splitmix64-style) used only for jitter -- avoids
96/// pulling in the `rand` crate for a test-only feature.
97fn next_pseudo_random() -> u64 {
98    static COUNTER: AtomicU64 = AtomicU64::new(0);
99    let counter = COUNTER.fetch_add(1, Ordering::Relaxed);
100    let nanos = SystemTime::now()
101        .duration_since(UNIX_EPOCH)
102        .map(|d| d.as_nanos() as u64)
103        .unwrap_or(0);
104    let mut x = nanos ^ counter.wrapping_mul(0x9E37_79B9_7F4A_7C15);
105    x ^= x >> 30;
106    x = x.wrapping_mul(0xBF58_476D_1CE4_E5B9);
107    x ^= x >> 27;
108    x = x.wrapping_mul(0x94D0_49BB_1331_11EB);
109    x ^= x >> 31;
110    x
111}
112
113#[derive(Clone, Debug, Default)]
114struct FaultState {
115    delay: [Option<Delay>; 2],
116    close_after: [Option<u64>; 2],
117    chunk_size: Option<usize>,
118    drop_all: bool,
119}
120
121/// A TCP proxy that forwards bytes between test clients and an upstream
122/// Redis node while injecting configurable network faults.
123///
124/// Construct with [`FaultProxy::spawn`]. Fault controls can be changed at
125/// any time via the returned handle; delay, close-after, and chunk-size take
126/// effect on data already in flight (checked on every read), while
127/// [`FaultProxy::set_drop_all`] is snapshotted once per accepted connection.
128/// The proxy stops accepting new connections when the handle is dropped;
129/// already-established connections run to completion.
130pub struct FaultProxy {
131    addr: SocketAddr,
132    state: Arc<RwLock<FaultState>>,
133    accept_task: JoinHandle<()>,
134}
135
136impl FaultProxy {
137    /// Bind an ephemeral local TCP listener that proxies to `upstream_addr`
138    /// and return a handle to it. The proxy starts in clean-passthrough mode.
139    pub async fn spawn(upstream_addr: impl ToSocketAddrs) -> Result<FaultProxy> {
140        let upstream_addr = tokio::net::lookup_host(upstream_addr)
141            .await?
142            .next()
143            .ok_or_else(|| {
144                crate::error::Error::Io(std::io::Error::new(
145                    std::io::ErrorKind::AddrNotAvailable,
146                    "could not resolve upstream address",
147                ))
148            })?;
149
150        let listener = TcpListener::bind(("127.0.0.1", 0)).await?;
151        let addr = listener.local_addr()?;
152        let state = Arc::new(RwLock::new(FaultState::default()));
153
154        let accept_state = state.clone();
155        let accept_task = tokio::spawn(async move {
156            loop {
157                let (client, _) = match listener.accept().await {
158                    Ok(pair) => pair,
159                    Err(_) => break,
160                };
161                tokio::spawn(handle_connection(
162                    client,
163                    upstream_addr,
164                    accept_state.clone(),
165                ));
166            }
167        });
168
169        Ok(FaultProxy {
170            addr,
171            state,
172            accept_task,
173        })
174    }
175
176    /// The local address clients should connect to.
177    pub fn addr(&self) -> SocketAddr {
178        self.addr
179    }
180
181    /// Set (or replace) the delay applied before forwarding data in `direction`.
182    pub fn set_delay(&self, direction: Direction, delay: Delay) {
183        self.state.write().unwrap().delay[direction.idx()] = Some(delay);
184    }
185
186    /// Remove any delay configured for `direction`.
187    pub fn clear_delay(&self, direction: Direction) {
188        self.state.write().unwrap().delay[direction.idx()] = None;
189    }
190
191    /// Close the connection once `bytes` total have been forwarded in
192    /// `direction`, cutting mid-buffer if the triggering read straddles the
193    /// threshold. Applies independently to every connection accepted while
194    /// this setting is active.
195    pub fn close_after(&self, direction: Direction, bytes: u64) {
196        self.state.write().unwrap().close_after[direction.idx()] = Some(bytes);
197    }
198
199    /// Remove the close-after threshold configured for `direction`.
200    pub fn clear_close_after(&self, direction: Direction) {
201        self.state.write().unwrap().close_after[direction.idx()] = None;
202    }
203
204    /// Split forwarded writes into pieces of at most `size` bytes, in both
205    /// directions. Use `size == 1` to exercise incremental decoders.
206    pub fn set_chunk_size(&self, size: usize) {
207        self.state.write().unwrap().chunk_size = Some(size.max(1));
208    }
209
210    /// Stop chunking writes; forward reads in whatever size they arrive.
211    pub fn clear_chunk_size(&self) {
212        self.state.write().unwrap().chunk_size = None;
213    }
214
215    /// Enable or disable black-hole mode. While enabled, newly accepted
216    /// connections are held open (no upstream connection is made, nothing is
217    /// ever forwarded or written back) until the client disconnects or the
218    /// proxy is dropped. Existing connections are unaffected.
219    pub fn set_drop_all(&self, drop_all: bool) {
220        self.state.write().unwrap().drop_all = drop_all;
221    }
222
223    /// Clear every fault control, returning the proxy to clean passthrough
224    /// for new connections.
225    pub fn reset(&self) {
226        *self.state.write().unwrap() = FaultState::default();
227    }
228}
229
230impl Drop for FaultProxy {
231    fn drop(&mut self) {
232        self.accept_task.abort();
233    }
234}
235
236async fn handle_connection(
237    mut client: TcpStream,
238    upstream_addr: SocketAddr,
239    state: Arc<RwLock<FaultState>>,
240) {
241    if state.read().unwrap().drop_all {
242        let mut buf = [0u8; 4096];
243        loop {
244            match client.read(&mut buf).await {
245                Ok(0) | Err(_) => return,
246                Ok(_) => {}
247            }
248        }
249    }
250
251    let upstream = match TcpStream::connect(upstream_addr).await {
252        Ok(s) => s,
253        Err(_) => return,
254    };
255
256    let (client_r, client_w) = client.into_split();
257    let (upstream_r, upstream_w) = upstream.into_split();
258
259    tokio::select! {
260        _ = forward(client_r, upstream_w, state.clone(), Direction::ClientToUpstream) => {},
261        _ = forward(upstream_r, client_w, state, Direction::UpstreamToClient) => {},
262    }
263}
264
265/// Copy bytes from `reader` to `writer`, applying whatever delay,
266/// close-after, and chunk-size controls are configured for `direction` at
267/// the time each chunk is read. Returns when the source hits EOF/error, the
268/// sink fails, or the close-after threshold for `direction` is reached.
269async fn forward(
270    mut reader: OwnedReadHalf,
271    mut writer: OwnedWriteHalf,
272    state: Arc<RwLock<FaultState>>,
273    direction: Direction,
274) {
275    let idx = direction.idx();
276    let mut total: u64 = 0;
277    let mut buf = [0u8; 4096];
278
279    loop {
280        let n = match reader.read(&mut buf).await {
281            Ok(0) | Err(_) => return,
282            Ok(n) => n,
283        };
284        let mut data = &buf[..n];
285
286        let (delay, close_after, chunk_size) = {
287            let s = state.read().unwrap();
288            (s.delay[idx], s.close_after[idx], s.chunk_size)
289        };
290
291        if let Some(limit) = close_after {
292            if total >= limit {
293                return;
294            }
295            let remaining = (limit - total) as usize;
296            if data.len() > remaining {
297                data = &data[..remaining];
298            }
299        }
300
301        if let Some(delay) = delay {
302            tokio::time::sleep(delay.resolve()).await;
303        }
304
305        let chunk = chunk_size.unwrap_or(data.len()).max(1);
306        for piece in data.chunks(chunk) {
307            if writer.write_all(piece).await.is_err() {
308                return;
309            }
310        }
311        total += data.len() as u64;
312
313        if let Some(limit) = close_after
314            && total >= limit
315        {
316            return;
317        }
318    }
319}
redis_server_wrapper/fault_proxy.rs

redis_server_wrapper/
fault_proxy.rs
