diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md index c5d6e27..a79b1e0 100644 --- a/ARCHITECTURE.md +++ b/ARCHITECTURE.md @@ -229,7 +229,7 @@ membership and failure detection use SWIM-style gossip. cluster configuration fo replication is a primary-to-replica stream built on top of the same persistence primitives used for local durability: -1. a replica connects and performs a handshake +1. a replica connects, and in a handshake each side proves it holds the cluster secret 2. the primary sends one in-memory snapshot per shard plus the current offset 3. the replica loads those snapshots and acknowledges their offsets 4. the primary streams incremental `AofRecord` updates, and the replica acknowledges them as it applies them diff --git a/CHANGELOG.md b/CHANGELOG.md index 8891d5c..d6811a7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -18,6 +18,12 @@ - cluster mode is about as fast as standalone without pipelining: about 860k requests a second against 320k on a 32-core machine. with a pipeline of 16 it serves about 2.3M against 320k, 15% below standalone. every command locked the migration state, an async mutex shared by every connection, and read the cluster state up to three times. it now reads each once, and the migration state takes a shared lock (#453) ### fixes +- replication hardening (#456): + - the replica checks that the primary holds the cluster secret before taking its data. only the primary checked the replica, so anything that could pose as the primary's address could replace a replica's data. each side signs the other's nonce after a label for its role, so one side's answer can't be relayed as the other's. the replication protocol version is now 3 + - the primary drops a connection that hasn't finished the handshake within 10 seconds, and runs at most 16 handshakes at once. an idle connection held a task forever before authenticating + - either side drops a connection on which a read or write makes no progress for 60 seconds, as Redis's `repl-timeout`. the primary sends a heartbeat after 10 seconds without data and the replica answers it. a primary that stopped responding left its replica waiting forever + - the replica reads snapshot and record bodies into buffers that grow as the bytes arrive. it allocated the full length from the header, up to 4 GiB, before reading any of it + - on a shard count mismatch the primary sends its status byte before closing. it was left unflushed, so the replica saw only a closed connection - WAIT counts a replica only once it has loaded its full sync. it counted the replica as soon as the primary had sent the snapshot, while the replica was still loading it, or had just flushed its data to start over (#455) - `INFO replication` on a primary counts the replicas streaming from it in `connected_replicas`. it counted the replicas gossip had reported, which could read 0 while one streamed. a new `sync_full` field counts the full syncs sent (#455) - failover safety (#452): diff --git a/crates/ember-server/src/replication.rs b/crates/ember-server/src/replication.rs index 0e9b9de..1470c11 100644 --- a/crates/ember-server/src/replication.rs +++ b/crates/ember-server/src/replication.rs @@ -11,16 +11,17 @@ //! //! ```text //! // Replica → primary (handshake request): -//! [version: 1B][num_shards: 2B] +//! [version: 1B][num_shards: 2B][replica_nonce: 16B] //! //! // Primary → replica (handshake response): //! [version: 1B][num_shards: 2B][primary_id_len: 1B][primary_id: N bytes] //! [status: 1B] (0 = ok, 1 = shard count mismatch) //! -//! // Authentication (if status = 0): -//! // primary → replica: [nonce: 16B] -//! // replica → primary: [tag: 32B] HMAC of the nonce with the cluster -//! // secret, or zeros without one +//! // Authentication (if status = 0). a tag is the HMAC of the role's +//! // label and the other side's nonce with the cluster secret, or zeros +//! // without one: +//! // primary → replica: [primary_nonce: 16B][primary tag: 32B] +//! // replica → primary: [replica tag: 32B] (sent once the primary's tag checks out) //! // primary → replica: [status: 1B] (0 = ok, 2 = authentication failed) //! //! // For each shard (if status = 0): @@ -36,16 +37,27 @@ //! //! // When the replica falls too far behind (see REPLICA_QUEUE_LIMIT): //! [MSG_RESYNC: 1B] primary closes the connection; replica reconnects +//! +//! // Primary → replica after HEARTBEAT_INTERVAL with nothing to send, and +//! // the replica's answer: +//! [MSG_PING: 1B] //! ``` //! +//! Either side drops a connection on which a read or write waits longer +//! than REPL_TIMEOUT, and the primary drops a replica that hasn't finished +//! the handshake within HANDSHAKE_TIMEOUT. +//! //! Offsets count replication records per shard. A snapshot's offset is the //! last record it contains, so the replica skips records at or below it. use std::collections::HashMap; +use std::future::Future; use std::io; use std::net::SocketAddr; +use std::pin::Pin; use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering}; use std::sync::{Arc, Mutex}; +use std::task::{Context, Poll}; use std::time::Duration; use bytes::{BufMut, Bytes, BytesMut}; @@ -53,15 +65,18 @@ use ember_cluster::ClusterSecret; use ember_core::{Engine, ShardRequest, ShardResponse}; use ember_persistence::aof::AofRecord; use ember_persistence::snapshot; -use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt, BufReader, BufWriter}; +use tokio::io::{ + AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt, BufReader, BufWriter, ReadBuf, +}; use tokio::net::{TcpListener, TcpStream}; -use tokio::sync::{broadcast, mpsc}; +use tokio::sync::{broadcast, mpsc, Semaphore}; +use tokio::time::{Instant, Sleep}; use tracing::{debug, error, info, warn}; // -- protocol constants -- -/// Version 2 sends real snapshot offsets and acknowledges `(shard, offset)`. -const REPL_VERSION: u8 = 2; +/// Version 3 authenticates both sides and adds heartbeats. +const REPL_VERSION: u8 = 3; const STATUS_OK: u8 = 0; const STATUS_SHARD_MISMATCH: u8 = 1; const STATUS_AUTH_FAILED: u8 = 2; @@ -72,6 +87,38 @@ const MSG_SHARD_OFFSET: u8 = 3; const MSG_RECORD: u8 = 4; const MSG_RESYNC: u8 = 5; const MSG_ACK: u8 = 6; +const MSG_PING: u8 = 7; + +/// Each side signs the other's nonce after a label for its own role, so +/// the answer one side gives can't be passed off as the other's. +const PRIMARY_ROLE: &[u8] = b"ember replication primary"; +const REPLICA_ROLE: &[u8] = b"ember replication replica"; + +/// How long the primary gives a connecting replica to finish the +/// handshake. Until then the peer is unauthenticated. Unit tests use +/// shorter times here and below. +const HANDSHAKE_TIMEOUT: Duration = test_or(Duration::from_secs(10), Duration::from_millis(200)); + +/// The most handshakes the primary runs at once. Connections beyond it +/// are closed. +const MAX_HANDSHAKES: usize = 16; + +/// How long either side waits on a read or write that makes no progress +/// before it drops the connection, as Redis's `repl-timeout`. +const REPL_TIMEOUT: Duration = test_or(Duration::from_secs(60), Duration::from_millis(600)); + +/// How long the primary goes without sending before it sends MSG_PING. +/// The replica answers with MSG_PING, so both directions carry data well +/// within [`REPL_TIMEOUT`]. +const HEARTBEAT_INTERVAL: Duration = test_or(Duration::from_secs(10), Duration::from_millis(100)); + +const fn test_or(normal: Duration, in_tests: Duration) -> Duration { + if cfg!(test) { + in_tests + } else { + normal + } +} /// The most record bytes the primary queues for one replica. Writes made /// while the replica loads its full sync wait in this queue. A replica @@ -180,6 +227,140 @@ impl ReplicaTracker { } } +/// The tag that proves the holder of `secret` speaks as `role`, answering +/// the other side's `nonce`. Without a secret it is all zeros. +fn auth_tag( + secret: Option<&ClusterSecret>, + role: &[u8], + nonce: &[u8; NONCE_LEN], +) -> [u8; ember_cluster::TAG_LEN] { + secret.map_or([0; ember_cluster::TAG_LEN], |secret| { + secret.sign(&[role, nonce].concat()) + }) +} + +/// Whether `tag` is [`auth_tag`] for `role` and `nonce`. Without a secret +/// any tag is accepted. +fn auth_ok( + secret: Option<&ClusterSecret>, + role: &[u8], + nonce: &[u8; NONCE_LEN], + tag: &[u8], +) -> bool { + secret.is_none_or(|secret| secret.verify(&[role, nonce].concat(), tag)) +} + +/// Runs `fut`, failing with `TimedOut` if it takes longer than `limit`. +async fn within( + limit: Duration, + what: &str, + fut: impl Future>, +) -> io::Result { + tokio::time::timeout(limit, fut).await.unwrap_or_else(|_| { + Err(io::Error::new( + io::ErrorKind::TimedOut, + format!("{what} timed out"), + )) + }) +} + +/// A connection whose reads and writes fail with `TimedOut` after waiting +/// [`REPL_TIMEOUT`] without progress. +/// +/// Replication connections stay open indefinitely, so a peer that stops +/// responding must not hold one forever. Heartbeats keep a quiet but +/// healthy connection busy. +struct IdleTimeout { + inner: S, + read: Deadline, + write: Deadline, +} + +type Conn = IdleTimeout; + +impl IdleTimeout { + fn new(inner: S) -> Self { + Self { + inner, + read: Deadline::new(), + write: Deadline::new(), + } + } +} + +/// The point at which one direction of an [`IdleTimeout`] gives up. +struct Deadline { + sleep: Pin>, + /// Set while an operation is waiting, from its first `Pending`. + armed: bool, +} + +impl Deadline { + fn new() -> Self { + Self { + sleep: Box::pin(tokio::time::sleep(Duration::ZERO)), + armed: false, + } + } + + /// Passes on the inner operation's `poll`, or fails it once it has + /// been pending for [`REPL_TIMEOUT`]. + fn check(&mut self, cx: &mut Context<'_>, poll: Poll>) -> Poll> { + if poll.is_ready() { + self.armed = false; + return poll; + } + if !self.armed { + self.sleep.as_mut().reset(Instant::now() + REPL_TIMEOUT); + self.armed = true; + } + match self.sleep.as_mut().poll(cx) { + Poll::Ready(()) => { + self.armed = false; + Poll::Ready(Err(io::Error::new( + io::ErrorKind::TimedOut, + "replication peer stopped responding", + ))) + } + Poll::Pending => Poll::Pending, + } + } +} + +impl AsyncRead for IdleTimeout { + fn poll_read( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &mut ReadBuf<'_>, + ) -> Poll> { + let this = self.get_mut(); + let poll = Pin::new(&mut this.inner).poll_read(cx, buf); + this.read.check(cx, poll) + } +} + +impl AsyncWrite for IdleTimeout { + fn poll_write( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &[u8], + ) -> Poll> { + let this = self.get_mut(); + let poll = Pin::new(&mut this.inner).poll_write(cx, buf); + this.write.check(cx, poll) + } + + fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let this = self.get_mut(); + let poll = Pin::new(&mut this.inner).poll_flush(cx); + this.write.check(cx, poll) + } + + fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + Pin::new(&mut self.get_mut().inner).poll_shutdown(cx) + } +} + // -- framed I/O primitives -- // // generic over AsyncRead/AsyncWrite so they work with both raw TcpStream @@ -227,6 +408,22 @@ async fn read_u64_le(r: &mut (impl AsyncRead + Unpin)) -> std::io::Result { Ok(u64::from_le_bytes(buf)) } +/// Reads a `len`-byte body into `buf`. The buffer grows as the bytes +/// arrive, so a corrupt length can't make the reader allocate memory the +/// sender never fills. +async fn read_body( + r: &mut (impl AsyncRead + Unpin), + len: u32, + buf: &mut Vec, +) -> io::Result<()> { + buf.clear(); + let read = r.take(u64::from(len)).read_to_end(buf).await?; + if read != len as usize { + return Err(io::ErrorKind::UnexpectedEof.into()); + } + Ok(()) +} + /// Reads one byte and returns an error if it doesn't match `expected`. /// /// Used throughout the handshake and snapshot-load path to assert that @@ -267,16 +464,17 @@ impl ReplicationServer { /// server bound to loopback does not expose its dataset on every /// interface. /// - /// Runs indefinitely in the background; returns immediately after - /// spawning the accept loop task. + /// Runs indefinitely in the background; returns the bound address + /// immediately after spawning the accept loop task. pub async fn start( engine: Arc, primary_id: String, addr: SocketAddr, tracker: Arc, secret: Option>, - ) -> std::io::Result<()> { + ) -> std::io::Result { let listener = TcpListener::bind(addr).await?; + let addr = listener.local_addr()?; info!(%addr, "replication server listening"); let server = Arc::new(Self { @@ -286,14 +484,19 @@ impl ReplicationServer { secret, }); + let handshakes = Arc::new(Semaphore::new(MAX_HANDSHAKES)); tokio::spawn(async move { loop { match listener.accept().await { Ok((stream, peer)) => { + let Ok(permit) = Arc::clone(&handshakes).try_acquire_owned() else { + warn!(%peer, "too many replication handshakes; closing connection"); + continue; + }; debug!(%peer, "replica connected"); let server = Arc::clone(&server); tokio::spawn(async move { - if let Err(e) = server.handle_replica(stream).await { + if let Err(e) = server.handle_replica(stream, permit).await { debug!(%peer, "replication connection closed: {e}"); } }); @@ -306,63 +509,27 @@ impl ReplicationServer { } }); - Ok(()) + Ok(addr) } /// Handles a single replica connection: handshake + full sync + stream. - async fn handle_replica(&self, stream: TcpStream) -> std::io::Result<()> { + /// `permit` holds one of the [`MAX_HANDSHAKES`] places until the + /// handshake ends. + async fn handle_replica( + &self, + stream: TcpStream, + permit: tokio::sync::OwnedSemaphorePermit, + ) -> std::io::Result<()> { // wrap in a 64 KiB write buffer so the 4-5 small framing writes // (tag, shard_id, offset, length) per record are coalesced into - // a single syscall. reads are delegated to the inner TcpStream. - let mut stream = BufWriter::with_capacity(65536, stream); - - // --- handshake --- - let replica_version = read_u8(&mut stream).await?; - if replica_version != REPL_VERSION { - return Err(std::io::Error::new( - std::io::ErrorKind::InvalidData, - format!("unsupported replication version: {replica_version}"), - )); - } - let replica_shards = read_u16_le(&mut stream).await?; - let our_shards = self.engine.shard_count() as u16; - - write_u8(&mut stream, REPL_VERSION).await?; - write_u16_le(&mut stream, our_shards).await?; - let id_bytes = self.primary_id.as_bytes(); - write_u8(&mut stream, id_bytes.len() as u8).await?; - stream.write_all(id_bytes).await?; - - if replica_shards != our_shards { - write_u8(&mut stream, STATUS_SHARD_MISMATCH).await?; - return Err(std::io::Error::new( - std::io::ErrorKind::InvalidData, - format!("shard count mismatch: replica={replica_shards} primary={our_shards}"), - )); - } - write_u8(&mut stream, STATUS_OK).await?; - - // the replica proves it holds the cluster secret before it gets - // any data. a fresh random nonce means a captured answer cannot be - // replayed. - let nonce: [u8; NONCE_LEN] = rand::random(); - stream.write_all(&nonce).await?; - stream.flush().await?; - let mut tag = [0u8; ember_cluster::TAG_LEN]; - stream.read_exact(&mut tag).await?; - let authenticated = self - .secret - .as_ref() - .is_none_or(|secret| secret.verify(&nonce, &tag)); - if !authenticated { - write_u8(&mut stream, STATUS_AUTH_FAILED).await?; - stream.flush().await?; - return Err(std::io::Error::new( - std::io::ErrorKind::PermissionDenied, - "replica failed cluster authentication", - )); - } - write_u8(&mut stream, STATUS_OK).await?; + // a single syscall. reads are delegated to the inner stream. + let mut stream = BufWriter::with_capacity(65536, IdleTimeout::new(stream)); + let handshake = within(HANDSHAKE_TIMEOUT, "handshake", self.handshake(&mut stream)).await; + // the status or error reply may still be in the buffer + let flushed = stream.flush().await; + handshake?; + flushed?; + drop(permit); // subscribe before snapshotting so we don't miss events that // happen between snapshot and incremental stream start @@ -461,6 +628,59 @@ impl ReplicationServer { result } + + /// Checks the replica's version and shard count, and that both sides + /// hold the cluster secret. + async fn handshake(&self, stream: &mut BufWriter) -> io::Result<()> { + let replica_version = read_u8(stream).await?; + if replica_version != REPL_VERSION { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!("unsupported replication version: {replica_version}"), + )); + } + let replica_shards = read_u16_le(stream).await?; + let mut replica_nonce = [0u8; NONCE_LEN]; + stream.read_exact(&mut replica_nonce).await?; + let our_shards = self.engine.shard_count() as u16; + + write_u8(stream, REPL_VERSION).await?; + write_u16_le(stream, our_shards).await?; + let id_bytes = self.primary_id.as_bytes(); + write_u8(stream, id_bytes.len() as u8).await?; + stream.write_all(id_bytes).await?; + + if replica_shards != our_shards { + write_u8(stream, STATUS_SHARD_MISMATCH).await?; + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!("shard count mismatch: replica={replica_shards} primary={our_shards}"), + )); + } + write_u8(stream, STATUS_OK).await?; + + // each side proves it holds the cluster secret by signing a fresh + // nonce from the other, so a captured answer can't be replayed. + // the primary answers first, and the replica gets no data until + // it has answered too + let secret = self.secret.as_deref(); + let nonce: [u8; NONCE_LEN] = rand::random(); + stream.write_all(&nonce).await?; + stream + .write_all(&auth_tag(secret, PRIMARY_ROLE, &replica_nonce)) + .await?; + stream.flush().await?; + let mut tag = [0u8; ember_cluster::TAG_LEN]; + stream.read_exact(&mut tag).await?; + if !auth_ok(secret, REPLICA_ROLE, &nonce, &tag) { + write_u8(stream, STATUS_AUTH_FAILED).await?; + return Err(io::Error::new( + io::ErrorKind::PermissionDenied, + "replica failed cluster authentication", + )); + } + write_u8(stream, STATUS_OK).await + } } /// Encodes the records of each replication event as MSG_RECORD frames and @@ -513,14 +733,25 @@ async fn forward_events( } /// Writes queued records to the replica, flushing whenever the queue runs -/// dry. Once the forwarder closes the queue, the rest of it is stale: it -/// sends MSG_RESYNC instead, so the replica starts over with a full sync. +/// dry, and MSG_PING when there has been nothing to send for +/// [`HEARTBEAT_INTERVAL`]. Once the forwarder closes the queue, the rest of +/// it is stale: it sends MSG_RESYNC instead, so the replica starts over +/// with a full sync. async fn stream_records( - writer: &mut BufWriter>, + writer: &mut BufWriter>, queue: &mut mpsc::UnboundedReceiver, queued: &AtomicUsize, ) -> std::io::Result<()> { - while let Some(frame) = queue.recv().await { + loop { + let frame = match tokio::time::timeout(HEARTBEAT_INTERVAL, queue.recv()).await { + Ok(Some(frame)) => frame, + Ok(None) => break, + Err(_) => { + write_u8(writer, MSG_PING).await?; + writer.flush().await?; + continue; + } + }; if queue.is_closed() { break; } @@ -602,53 +833,8 @@ impl ReplicationClient { /// fine for hours does not wait the maximum delay after one drop. async fn sync(&self, stream: TcpStream, backoff: &mut Duration) -> std::io::Result<()> { // reads come through a buffer; writes go straight to the socket - let mut stream = BufReader::with_capacity(65536, stream); - let our_shards = self.engine.shard_count() as u16; - - // send handshake - write_u8(&mut stream, REPL_VERSION).await?; - write_u16_le(&mut stream, our_shards).await?; - - // read primary response - let primary_version = read_u8(&mut stream).await?; - if primary_version != REPL_VERSION { - return Err(std::io::Error::new( - std::io::ErrorKind::InvalidData, - format!("unsupported primary replication version: {primary_version}"), - )); - } - let primary_shards = read_u16_le(&mut stream).await?; - let id_len = read_u8(&mut stream).await? as usize; - let mut id_buf = vec![0u8; id_len]; - stream.read_exact(&mut id_buf).await?; - let primary_id = String::from_utf8_lossy(&id_buf).into_owned(); - - let status = read_u8(&mut stream).await?; - if status == STATUS_SHARD_MISMATCH { - return Err(std::io::Error::new( - std::io::ErrorKind::InvalidData, - format!( - "shard count mismatch with primary {primary_id}: \ - ours={our_shards} primary={primary_shards}" - ), - )); - } - - // answer the primary's challenge with the cluster secret - let mut nonce = [0u8; NONCE_LEN]; - stream.read_exact(&mut nonce).await?; - let tag = self - .secret - .as_ref() - .map(|secret| secret.sign(&nonce)) - .unwrap_or([0u8; ember_cluster::TAG_LEN]); - stream.write_all(&tag).await?; - if read_u8(&mut stream).await? == STATUS_AUTH_FAILED { - return Err(std::io::Error::new( - std::io::ErrorKind::PermissionDenied, - format!("primary {primary_id} rejected our cluster credentials"), - )); - } + let mut stream = BufReader::with_capacity(65536, IdleTimeout::new(stream)); + let (primary_id, primary_shards) = self.handshake(&mut stream).await?; info!(primary_id = %primary_id, "handshake ok, loading full sync"); *backoff = INITIAL_BACKOFF; @@ -667,11 +853,8 @@ impl ReplicationClient { for _ in 0..primary_shards { expect_tag(&mut stream, MSG_SHARD_SYNC, "MSG_SHARD_SYNC").await?; let shard_id = read_u16_le(&mut stream).await?; - let snap_len = read_u32_le(&mut stream).await? as usize; - - snap_buf.clear(); - snap_buf.resize(snap_len, 0); - stream.read_exact(&mut snap_buf).await?; + let snap_len = read_u32_le(&mut stream).await?; + read_body(&mut stream, snap_len, &mut snap_buf).await?; self.apply_snapshot(&mut batch, shard_id, &snap_buf).await?; expect_tag(&mut stream, MSG_SHARD_OFFSET, "MSG_SHARD_OFFSET").await?; @@ -694,15 +877,15 @@ impl ReplicationClient { // incremental stream. records are collected into a batch, which // is applied and acknowledged once the buffered input is used up + let mut record_bytes = Vec::new(); loop { let msg = read_u8(&mut stream).await?; match msg { MSG_RECORD => { let shard_id = read_u16_le(&mut stream).await?; let offset = read_u64_le(&mut stream).await?; - let record_len = read_u32_le(&mut stream).await? as usize; - let mut record_bytes = vec![0u8; record_len]; - stream.read_exact(&mut record_bytes).await?; + let record_len = read_u32_le(&mut stream).await?; + read_body(&mut stream, record_len, &mut record_bytes).await?; let record = AofRecord::from_bytes(&record_bytes).map_err(|e| { std::io::Error::new( @@ -726,6 +909,7 @@ impl ReplicationClient { .flatten(); batch.push(shard_id, request, Some(offset)); } + MSG_PING => write_u8(&mut stream, MSG_PING).await?, MSG_RESYNC => { info!("primary requested resync; reconnecting"); return Ok(()); @@ -752,6 +936,63 @@ impl ReplicationClient { } } + /// Checks the primary's version and shard count, and that both sides + /// hold the cluster secret. Returns the primary's id and shard count. + async fn handshake(&self, stream: &mut BufReader) -> io::Result<(String, u16)> { + let our_shards = self.engine.shard_count() as u16; + let nonce: [u8; NONCE_LEN] = rand::random(); + write_u8(stream, REPL_VERSION).await?; + write_u16_le(stream, our_shards).await?; + stream.write_all(&nonce).await?; + + let primary_version = read_u8(stream).await?; + if primary_version != REPL_VERSION { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!("unsupported primary replication version: {primary_version}"), + )); + } + let primary_shards = read_u16_le(stream).await?; + let id_len = read_u8(stream).await? as usize; + let mut id_buf = vec![0u8; id_len]; + stream.read_exact(&mut id_buf).await?; + let primary_id = String::from_utf8_lossy(&id_buf).into_owned(); + + if read_u8(stream).await? == STATUS_SHARD_MISMATCH { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!( + "shard count mismatch with primary {primary_id}: \ + ours={our_shards} primary={primary_shards}" + ), + )); + } + + // the primary proves it holds the cluster secret before this node + // flushes its data for the primary's, then this node proves it + let secret = self.secret.as_deref(); + let mut primary_nonce = [0u8; NONCE_LEN]; + stream.read_exact(&mut primary_nonce).await?; + let mut tag = [0u8; ember_cluster::TAG_LEN]; + stream.read_exact(&mut tag).await?; + if !auth_ok(secret, PRIMARY_ROLE, &nonce, &tag) { + return Err(io::Error::new( + io::ErrorKind::PermissionDenied, + format!("primary {primary_id} failed cluster authentication"), + )); + } + stream + .write_all(&auth_tag(secret, REPLICA_ROLE, &primary_nonce)) + .await?; + if read_u8(stream).await? == STATUS_AUTH_FAILED { + return Err(io::Error::new( + io::ErrorKind::PermissionDenied, + format!("primary {primary_id} rejected our cluster credentials"), + )); + } + Ok((primary_id, primary_shards)) + } + /// Loads a snapshot blob into the given shard. /// /// Each entry goes in through RESTORE, which replaces the key with the @@ -867,6 +1108,162 @@ impl Batch { #[cfg(test)] mod tests { use super::*; + use ember_core::{EngineConfig, Value}; + + fn secret(password: &str) -> Option> { + Some(Arc::new(ClusterSecret::from_password(password))) + } + + fn set(key: &str, value: &'static str) -> ShardRequest { + ShardRequest::Set { + key: key.into(), + value: Bytes::from_static(value.as_bytes()), + expire: None, + nx: false, + xx: false, + } + } + + /// The value of `key` on `engine`, if it is a string. + async fn get(engine: &Engine, key: &str) -> Option { + let get = ShardRequest::Get { key: key.into() }; + match engine.send_to_shard(0, get).await.unwrap() { + ShardResponse::Value(Some(Value::String(v))) => Some(v), + _ => None, + } + } + + /// Starts a primary with one shard holding `greeting = hello`. Returns + /// its replication address and engine. + async fn start_primary(secret: Option>) -> (SocketAddr, Arc) { + let (tx, _) = broadcast::channel(16); + let config = EngineConfig { + replication_tx: Some(tx), + ..Default::default() + }; + let engine = Arc::new(Engine::with_config(1, config)); + engine + .send_to_shard(0, set("greeting", "hello")) + .await + .unwrap(); + let addr = "127.0.0.1:0".parse().unwrap(); + let tracker = Arc::new(ReplicaTracker::new()); + let addr = + ReplicationServer::start(Arc::clone(&engine), "primary".into(), addr, tracker, secret) + .await + .unwrap(); + (addr, engine) + } + + fn client(secret: Option>) -> ReplicationClient { + ReplicationClient { + engine: Arc::new(Engine::new(1)), + primary_addr: "127.0.0.1:1".parse().unwrap(), + secret, + } + } + + /// Runs one replication session against `addr`, for at most two + /// replication timeouts. + async fn sync_once(client: &ReplicationClient, addr: SocketAddr) -> io::Result<()> { + let stream = TcpStream::connect(addr).await?; + let mut backoff = INITIAL_BACKOFF; + within(REPL_TIMEOUT * 2, "sync", client.sync(stream, &mut backoff)).await + } + + #[tokio::test] + async fn each_side_checks_that_the_other_holds_the_secret() { + let denied = |result: io::Result<()>| matches!(result, Err(e) if e.kind() == io::ErrorKind::PermissionDenied); + + // a primary with another secret, or none, can't feed this replica + let (primary, _) = start_primary(secret("other")).await; + assert!(denied(sync_once(&client(secret("ours")), primary).await)); + let (open_primary, _) = start_primary(None).await; + assert!(denied( + sync_once(&client(secret("ours")), open_primary).await + )); + + // a replica without the secret gets nothing + let (primary, _) = start_primary(secret("ours")).await; + assert!(denied(sync_once(&client(None), primary).await)); + + // with the same secret the replica syncs, then streams until the + // test stops it + let replica = client(secret("ours")); + let result = sync_once(&replica, primary).await; + assert!(matches!(result, Err(e) if e.kind() == io::ErrorKind::TimedOut)); + assert_eq!(get(&replica.engine, "greeting").await.unwrap(), "hello"); + } + + #[tokio::test] + async fn heartbeats_keep_a_quiet_connection_open() { + let (primary, primary_engine) = start_primary(None).await; + let replica = Arc::new(client(None)); + let session = tokio::spawn({ + let replica = Arc::clone(&replica); + async move { + let stream = TcpStream::connect(primary).await?; + let mut backoff = INITIAL_BACKOFF; + replica.sync(stream, &mut backoff).await + } + }); + + tokio::time::sleep(REPL_TIMEOUT * 3).await; + assert!(!session.is_finished(), "{:?}", session.await); + + // the connection still carries writes + primary_engine + .send_to_shard(0, set("later", "still here")) + .await + .unwrap(); + for _ in 0..100 { + if get(&replica.engine, "later").await.is_some() { + return; + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + panic!("the write never reached the replica"); + } + + #[tokio::test] + async fn the_primary_drops_a_peer_that_never_finishes_the_handshake() { + let (primary, _) = start_primary(None).await; + let start = Instant::now(); + let mut peer = TcpStream::connect(primary).await.unwrap(); + // the primary closes the connection without a byte + let mut buf = [0u8; 1]; + assert_eq!(peer.read(&mut buf).await.unwrap(), 0); + assert!(start.elapsed() >= HANDSHAKE_TIMEOUT); + } + + #[tokio::test] + async fn a_replica_drops_a_primary_that_goes_quiet() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + // a primary with no shards that completes the handshake, then + // sends nothing, not even heartbeats + tokio::spawn(async move { + let (mut stream, _) = listener.accept().await.unwrap(); + let mut request = [0u8; 3 + NONCE_LEN]; + stream.read_exact(&mut request).await.unwrap(); + let mut reply = vec![REPL_VERSION, 0, 0, 1, b'p', STATUS_OK]; + reply.extend_from_slice(&[0; NONCE_LEN]); + reply.extend_from_slice(&[0; ember_cluster::TAG_LEN]); + stream.write_all(&reply).await.unwrap(); + let mut tag = [0u8; ember_cluster::TAG_LEN]; + stream.read_exact(&mut tag).await.unwrap(); + stream.write_all(&[STATUS_OK]).await.unwrap(); + std::future::pending::<()>().await; + }); + + let replica = client(None); + let start = Instant::now(); + let stream = TcpStream::connect(addr).await.unwrap(); + let mut backoff = INITIAL_BACKOFF; + let result = replica.sync(stream, &mut backoff).await; + assert!(matches!(result, Err(e) if e.kind() == io::ErrorKind::TimedOut)); + assert!(start.elapsed() >= REPL_TIMEOUT); + } #[test] fn tracker_counts_replicas_that_reached_every_shard() { diff --git a/docs/production-checklist.md b/docs/production-checklist.md index c46c22a..c9bb45d 100644 --- a/docs/production-checklist.md +++ b/docs/production-checklist.md @@ -264,7 +264,7 @@ in clustered deployments, gossip and replication traffic between nodes is authen auth-pass = "shared-cluster-secret" ``` -nodes that do not present the correct HMAC are rejected, and a replica without the key cannot pull data from its primary. a cluster node bound to anything other than loopback refuses to start without this setting. +nodes that do not present the correct HMAC are rejected. a replication connection checks both ends: a replica without the key cannot pull data from its primary, and a replica takes no data from a primary without the key. only the replication handshake is authenticated, not each record after it, and the stream is not encrypted, so keep replication traffic on a trusted network. a cluster node bound to anything other than loopback refuses to start without this setting. **key and value size limits**