dynomite/entropy/
send.rs

1//! Entropy sender.
2//!
3//! Drives the client side of the reconciliation channel: dial the
4//! peer, negotiate, then push a snapshot from a [`SnapshotSource`]
5//! one chunk at a time.
6
7use std::io::{Read, Write};
8use std::net::{SocketAddr, TcpStream};
9use std::time::Duration;
10
11use aes::Aes128;
12use cbc::cipher::block_padding::Pkcs7;
13use cbc::cipher::{BlockEncryptMut, KeyIvInit};
14use tokio::io::AsyncWriteExt;
15use tokio::net::TcpStream as TokioTcpStream;
16use tokio::task::JoinHandle;
17
18use crate::entropy::util::EntropyMaterial;
19use crate::entropy::{
20    BoxedSnapshotSource, EntropyConfig, EntropyError, EntropyResult, NegotiationHeader,
21    SnapshotHeader, SnapshotSource, ENTROPY_COMMAND_SEND, ENTROPY_MAGIC,
22};
23
24type Aes128CbcEnc = cbc::Encryptor<Aes128>;
25
26/// Entropy sender driver.
27///
28/// Operationally, calling [`run`](EntropySender::run) spawns a
29/// tokio task that performs one push and then exits.
30///
31/// # Examples
32///
33/// ```no_run
34/// use std::path::PathBuf;
35/// use std::sync::Arc;
36/// use dynomite::entropy::{
37///     EntropyConfig, EntropySender, RedisLocalSnapshot,
38/// };
39///
40/// # async fn run() {
41/// let cfg = EntropyConfig {
42///     key_file: PathBuf::from("conf/recon_key.pem"),
43///     iv_file: PathBuf::from("conf/recon_iv.pem"),
44///     listen_addr: "127.0.0.1:8105".parse().unwrap(),
45///     send_addr: None,
46///     peer_endpoint: "127.0.0.1:8105".parse().unwrap(),
47///     buffer_size: 16 * 1024,
48///     header_size: 1024,
49///     encrypt: true,
50/// };
51/// let source = Arc::new(RedisLocalSnapshot::default());
52/// let handle = EntropySender::run(cfg, source);
53/// handle.await.unwrap().unwrap();
54/// # }
55/// ```
56pub struct EntropySender;
57
58impl EntropySender {
59    /// Spawn a tokio task that performs one snapshot push.
60    ///
61    /// The returned [`JoinHandle`] resolves to the push result.
62    /// The task uses `cfg.peer_endpoint` to dial the receiver and
63    /// `cfg.send_addr` (when set) for the local bind.
64    #[must_use]
65    pub fn run(
66        cfg: EntropyConfig,
67        source: BoxedSnapshotSource,
68    ) -> JoinHandle<EntropyResult<usize>> {
69        tokio::spawn(async move { Self::push(cfg, source).await })
70    }
71
72    /// Perform one snapshot push and return the number of plaintext
73    /// bytes transferred. Used directly by tests; the spawned
74    /// variant just calls this on a tokio task.
75    ///
76    /// # Errors
77    /// Forwards anything from key loading, source acquisition,
78    /// dialing, or transport.
79    pub async fn push(cfg: EntropyConfig, source: BoxedSnapshotSource) -> EntropyResult<usize> {
80        Self::push_with_material(cfg, source, None).await
81    }
82
83    /// Perform one snapshot push using a pre-loaded
84    /// [`EntropyMaterial`].
85    ///
86    /// When `override_material` is `None` and `cfg.encrypt` is
87    /// set, key + IV bytes are loaded from `cfg.key_file` /
88    /// `cfg.iv_file` (the historical behaviour). When
89    /// `override_material` is `Some(_)`, the supplied material is
90    /// used verbatim and the on-disk paths are not touched. The
91    /// override is consumed by the [`crate::entropy::driver::EntropyDriver`]
92    /// run loop, which loads the AES material once at startup
93    /// and drives many cycles from the same handle.
94    ///
95    /// # Errors
96    /// Forwards anything from key loading (when no override is
97    /// supplied), source acquisition, dialing, or transport.
98    pub async fn push_with_material(
99        cfg: EntropyConfig,
100        source: BoxedSnapshotSource,
101        override_material: Option<EntropyMaterial>,
102    ) -> EntropyResult<usize> {
103        cfg.validate()?;
104        let material = if cfg.encrypt {
105            match override_material {
106                Some(m) => Some(m),
107                None => Some(crate::entropy::util::load_material(
108                    &cfg.key_file,
109                    &cfg.iv_file,
110                )?),
111            }
112        } else {
113            None
114        };
115        let snapshot = collect_snapshot(source).await?;
116
117        let mut stream = dial(&cfg).await?;
118        write_negotiation(&mut stream, &cfg).await?;
119        write_snapshot_header(&mut stream, &cfg, snapshot.len()).await?;
120        write_chunks(&mut stream, &cfg, &snapshot, material.as_ref()).await?;
121        stream.shutdown().await?;
122        Ok(snapshot.len())
123    }
124}
125
126async fn collect_snapshot(source: BoxedSnapshotSource) -> EntropyResult<Vec<u8>> {
127    tokio::task::spawn_blocking(move || source.snapshot())
128        .await
129        .map_err(|e| EntropyError::Source(format!("snapshot task panicked: {e}")))?
130}
131
132async fn dial(cfg: &EntropyConfig) -> EntropyResult<TokioTcpStream> {
133    let socket = match cfg.peer_endpoint {
134        SocketAddr::V4(_) => tokio::net::TcpSocket::new_v4()?,
135        SocketAddr::V6(_) => tokio::net::TcpSocket::new_v6()?,
136    };
137    if let Some(local) = cfg.send_addr {
138        socket.bind(local)?;
139    }
140    let stream = socket.connect(cfg.peer_endpoint).await?;
141    stream.set_nodelay(true)?;
142    Ok(stream)
143}
144
145async fn write_negotiation(stream: &mut TokioTcpStream, cfg: &EntropyConfig) -> EntropyResult<()> {
146    let hdr = NegotiationHeader {
147        magic: ENTROPY_MAGIC,
148        command: ENTROPY_COMMAND_SEND,
149        header_size: u32::try_from(cfg.header_size)
150            .map_err(|_| EntropyError::Config("header_size > u32::MAX".to_string()))?,
151        buffer_size: u32::try_from(cfg.buffer_size)
152            .map_err(|_| EntropyError::Config("buffer_size > u32::MAX".to_string()))?,
153        cipher_size: u32::try_from(cipher_capacity(cfg.buffer_size))
154            .map_err(|_| EntropyError::Config("cipher_size > u32::MAX".to_string()))?,
155    };
156    stream.write_all(&hdr.to_wire()).await?;
157    Ok(())
158}
159
160async fn write_snapshot_header(
161    stream: &mut TokioTcpStream,
162    cfg: &EntropyConfig,
163    total_len: usize,
164) -> EntropyResult<()> {
165    let total_len_u32 = u32::try_from(total_len)
166        .map_err(|_| EntropyError::Source(format!("snapshot too large: {total_len} bytes")))?;
167    let hdr = SnapshotHeader {
168        total_len: total_len_u32,
169        encrypt_flag: u32::from(cfg.encrypt),
170    };
171    let bytes = hdr.to_wire(cfg.header_size)?;
172    stream.write_all(&bytes).await?;
173    Ok(())
174}
175
176async fn write_chunks(
177    stream: &mut TokioTcpStream,
178    cfg: &EntropyConfig,
179    snapshot: &[u8],
180    material: Option<&EntropyMaterial>,
181) -> EntropyResult<()> {
182    let buf = cfg.buffer_size;
183    let mut offset = 0;
184    while offset < snapshot.len() {
185        let end = (offset + buf).min(snapshot.len());
186        let plaintext = &snapshot[offset..end];
187        let payload: Vec<u8> = if let Some(mat) = material {
188            encrypt_chunk(plaintext, mat)?
189        } else {
190            plaintext.to_vec()
191        };
192        let chunk_len = u32::try_from(payload.len())
193            .map_err(|_| EntropyError::Protocol("chunk too large".to_string()))?;
194        stream.write_all(&chunk_len.to_be_bytes()).await?;
195        stream.write_all(&payload).await?;
196        offset = end;
197    }
198    Ok(())
199}
200
201/// AES-128-CBC encrypt one chunk using the entropy key + IV.
202///
203/// Pure helper exposed for the integration tests.
204///
205/// # Errors
206/// [`EntropyError::Crypto`] if the cipher reports a fault.
207pub fn encrypt_chunk(plaintext: &[u8], material: &EntropyMaterial) -> EntropyResult<Vec<u8>> {
208    let key = material.key().as_bytes();
209    let iv = material.iv().as_bytes();
210    let cipher = Aes128CbcEnc::new(key.into(), iv.into());
211    Ok(cipher.encrypt_padded_vec_mut::<Pkcs7>(plaintext))
212}
213
214/// PKCS#7 cipher overhead is at most one block; the negotiated
215/// cipher capacity is therefore plaintext + 16 bytes. Production
216/// servers can safely advertise more headroom.
217pub(crate) fn cipher_capacity(buffer_size: usize) -> usize {
218    buffer_size + 16
219}
220
221/// Default snapshot source: pulls a Redis snapshot from a local
222/// Redis instance.
223///
224/// The default source issues a `BGREWRITEAOF` over the
225/// configured Redis TCP endpoint, waits for the command to be
226/// acknowledged, then reads the on-disk AOF file from disk and
227/// pipes it over the entropy channel. Embedders
228/// that need a different snapshot strategy plug in their own
229/// [`SnapshotSource`] through the embedding API.
230///
231/// # Examples
232///
233/// ```
234/// use dynomite::entropy::RedisLocalSnapshot;
235/// let source = RedisLocalSnapshot::default();
236/// drop(source);
237/// ```
238pub struct RedisLocalSnapshot {
239    /// Address of the local Redis instance.
240    pub redis_addr: SocketAddr,
241    /// Path to the AOF file on disk.
242    pub aof_path: std::path::PathBuf,
243    /// Connect/read timeout.
244    pub timeout: Duration,
245    /// Number of `BGREWRITEAOF` retries before giving up. The
246    /// count and the retry pause are exposed as parameters.
247    pub bgrewrite_retries: u32,
248    /// Pause between `BGREWRITEAOF` retries.
249    pub bgrewrite_retry_pause: Duration,
250}
251
252impl Default for RedisLocalSnapshot {
253    fn default() -> Self {
254        Self {
255            redis_addr: "127.0.0.1:22122".parse().expect("static literal parses"),
256            aof_path: std::path::PathBuf::from("/mnt/data/nfredis/appendonly.aof"),
257            timeout: Duration::from_secs(30),
258            bgrewrite_retries: 1,
259            bgrewrite_retry_pause: Duration::from_secs(10),
260        }
261    }
262}
263
264impl RedisLocalSnapshot {
265    /// Override the Redis address.
266    #[must_use]
267    pub fn with_redis_addr(mut self, addr: SocketAddr) -> Self {
268        self.redis_addr = addr;
269        self
270    }
271
272    /// Override the AOF on-disk path.
273    #[must_use]
274    pub fn with_aof_path(mut self, path: std::path::PathBuf) -> Self {
275        self.aof_path = path;
276        self
277    }
278
279    fn bgrewriteaof(&self) -> EntropyResult<()> {
280        let mut last_err: Option<EntropyError> = None;
281        for attempt in 0..=self.bgrewrite_retries {
282            match self.try_bgrewriteaof() {
283                Ok(()) => return Ok(()),
284                Err(e) => {
285                    last_err = Some(e);
286                    if attempt < self.bgrewrite_retries {
287                        std::thread::sleep(self.bgrewrite_retry_pause);
288                    }
289                }
290            }
291        }
292        Err(last_err.unwrap_or_else(|| {
293            EntropyError::Source("bgrewriteaof failed without error".to_string())
294        }))
295    }
296
297    fn try_bgrewriteaof(&self) -> EntropyResult<()> {
298        let mut sock = TcpStream::connect_timeout(&self.redis_addr, self.timeout)
299            .map_err(|e| EntropyError::Source(format!("connect to redis: {e}")))?;
300        sock.set_read_timeout(Some(self.timeout))
301            .map_err(|e| EntropyError::Source(format!("redis timeout: {e}")))?;
302        sock.set_write_timeout(Some(self.timeout))
303            .map_err(|e| EntropyError::Source(format!("redis timeout: {e}")))?;
304        sock.write_all(b"*1\r\n$13\r\nBGREWRITEAOF\r\n")
305            .map_err(|e| EntropyError::Source(format!("redis write: {e}")))?;
306        let mut buf = [0u8; 256];
307        let n = sock
308            .read(&mut buf)
309            .map_err(|e| EntropyError::Source(format!("redis read: {e}")))?;
310        let reply = &buf[..n];
311        if reply.first() != Some(&b'+') {
312            return Err(EntropyError::Source(format!(
313                "BGREWRITEAOF rejected: {}",
314                String::from_utf8_lossy(reply)
315            )));
316        }
317        Ok(())
318    }
319}
320
321impl SnapshotSource for RedisLocalSnapshot {
322    fn snapshot(&self) -> EntropyResult<Vec<u8>> {
323        self.bgrewriteaof()?;
324        // Brief pause to let Redis flush the AOF rewrite to disk.
325        std::thread::sleep(Duration::from_secs(1));
326        let bytes = std::fs::read(&self.aof_path).map_err(|e| {
327            EntropyError::Source(format!("read AOF {}: {e}", self.aof_path.display()))
328        })?;
329        Ok(bytes)
330    }
331}
332
333/// In-memory snapshot source useful in tests and embedders that
334/// already hold the snapshot in RAM.
335///
336/// # Examples
337///
338/// ```
339/// use std::sync::Arc;
340/// use dynomite::entropy::send::StaticSnapshot;
341/// use dynomite::entropy::SnapshotSource;
342/// let s = Arc::new(StaticSnapshot::new(b"hello".to_vec()));
343/// assert_eq!(s.snapshot().unwrap(), b"hello");
344/// ```
345pub struct StaticSnapshot {
346    bytes: Vec<u8>,
347}
348
349impl StaticSnapshot {
350    /// Wrap a fixed byte vector.
351    #[must_use]
352    pub fn new(bytes: Vec<u8>) -> Self {
353        Self { bytes }
354    }
355}
356
357impl SnapshotSource for StaticSnapshot {
358    fn snapshot(&self) -> EntropyResult<Vec<u8>> {
359        Ok(self.bytes.clone())
360    }
361}
362
363#[cfg(test)]
364mod tests {
365    use super::*;
366    use crate::entropy::util::{EntropyIv, EntropyKey, ENTROPY_IV_LEN, ENTROPY_KEY_LEN};
367
368    fn material() -> EntropyMaterial {
369        EntropyMaterial::new(
370            EntropyKey::from_bytes([0x10; ENTROPY_KEY_LEN]),
371            EntropyIv::from_bytes([0x42; ENTROPY_IV_LEN]),
372        )
373    }
374
375    #[test]
376    fn encrypt_chunk_round_trips_with_pkcs7() {
377        use aes::Aes128;
378        use cbc::cipher::block_padding::Pkcs7;
379        use cbc::cipher::{BlockDecryptMut, KeyIvInit};
380        type Dec = cbc::Decryptor<Aes128>;
381
382        let mat = material();
383        let pt = b"hello entropy world";
384        let ct = encrypt_chunk(pt, &mat).unwrap();
385        // PKCS#7 always extends to the next 16-byte block.
386        assert!(ct.len() >= pt.len());
387        assert_eq!(ct.len() % 16, 0);
388
389        let dec = Dec::new(mat.key().as_bytes().into(), mat.iv().as_bytes().into());
390        let plain = dec.decrypt_padded_vec_mut::<Pkcs7>(&ct).unwrap();
391        assert_eq!(plain, pt);
392    }
393
394    #[test]
395    fn cipher_capacity_includes_pkcs7_block() {
396        assert_eq!(cipher_capacity(16), 32);
397        assert_eq!(cipher_capacity(15), 31);
398    }
399
400    #[test]
401    fn static_snapshot_returns_payload() {
402        let s = StaticSnapshot::new(b"abc".to_vec());
403        assert_eq!(s.snapshot().unwrap(), b"abc");
404    }
405
406    #[test]
407    fn arc_static_snapshot_returns_payload() {
408        let s: BoxedSnapshotSource = std::sync::Arc::new(StaticSnapshot::new(vec![1, 2, 3]));
409        assert_eq!(s.snapshot().unwrap(), vec![1, 2, 3]);
410    }
411}
dynomite/entropy/send.rs

dynomite/entropy/
send.rs
