fix: close connection, auth and transaction gaps

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
+- connections and auth (#458):
+  - MULTI, EXEC and WATCH work in a batch that also holds AUTH. such a batch skipped transaction handling: MULTI was refused, its commands ran at once, and an open transaction stopped queueing
+  - ACL SETUSER applies all of its rules or, when one is invalid, none. it kept the rules before the invalid one, and created the user
+  - a transaction's queued commands and watched keys may hold at most 128 MiB, the most one command may. the queue had no limit. a command past it is refused, and EXEC then discards the transaction
+  - CONFIG REWRITE keeps the config file's permissions, and writes it readable by the owner only until then. it wrote the file, which can hold requirepass, with the default permissions, often readable by everyone
+  - gRPC connections count toward `maxclients`
 - gRPC uses TLS when `tls-port` is set, with the same certificate, key and client certificate settings. it served plaintext on its own port even then, so a server that encrypted RESP still sent the same data in the clear over gRPC. the server refuses to start when the TLS port and the gRPC port are the same (#457)
 - 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

crates/ember-server/src/acl.rs

@@ -425,15 +425,14 @@ impl AclState {
         self.users.get(name).cloned().map(Arc::new)
     }
 
-    /// Creates or updates a user by applying rules.
+    /// Creates or updates a user by applying rules. The rules apply all
+    /// together or, if one is invalid, not at all.
     pub fn set_user(&mut self, name: &str, rules: &[String]) -> Result<(), String> {
-        let user = self
-            .users
-            .entry(name.to_string())
-            .or_insert_with(AclUser::new);
+        let mut user = self.users.get(name).cloned().unwrap_or_else(AclUser::new);
         for rule in rules {
-            apply_rule(user, rule)?;
+            apply_rule(&mut user, rule)?;
         }
+        self.users.insert(name.to_string(), user);
         Ok(())
     }
 
@@ -1140,6 +1139,21 @@ mod tests {
         assert!(err.contains("bogus"));
     }
 
+    #[test]
+    fn setuser_applies_no_rule_when_one_is_invalid() {
+        let mut acl = AclState::new(None);
+        let rules = |r: &[&str]| r.iter().map(|s| s.to_string()).collect::<Vec<_>>();
+        acl.set_user("alice", &rules(&["on", ">pw"])).unwrap();
+
+        assert!(acl.set_user("alice", &rules(&["off", "bogus"])).is_err());
+        assert!(
+            acl.get_user("alice").unwrap().enabled,
+            "the valid rule applied"
+        );
+        assert!(acl.set_user("bob", &rules(&["on", "bogus"])).is_err());
+        assert!(acl.get_user("bob").is_none(), "the user was created");
+    }
+
     #[test]
     fn setuser_and_list() {
         let mut state = AclState::new(None);

crates/ember-server/src/config.rs

@@ -779,21 +779,52 @@ impl ConfigRegistry {
 
         let toml_str = cfg.to_toml()?;
 
-        // atomic write: write to temp file, then rename
-        let tmp_path = config_path.with_extension("tmp");
-        std::fs::write(&tmp_path, toml_str.as_bytes())
-            .map_err(|e| format!("ERR failed to write config: {e}"))?;
-        std::fs::rename(&tmp_path, config_path)
-            .map_err(|e| format!("ERR failed to rename config file: {e}"))?;
-
-        Ok(())
+        write_private(config_path, toml_str.as_bytes())
+            .map_err(|e| format!("ERR failed to write config: {e}"))
     }
 }
 
+/// Replaces `path` with `contents` atomically, through a temp file and a
+/// rename. The file may hold requirepass, so the temp file is created
+/// readable by the owner only, then given the permissions `path` had.
+fn write_private(path: &Path, contents: &[u8]) -> std::io::Result<()> {
+    use std::io::Write;
+
+    let permissions = std::fs::metadata(path)?.permissions();
+    let tmp_path = path.with_extension("tmp");
+    // a temp file left by a crash keeps its old permissions when reopened
+    let _ = std::fs::remove_file(&tmp_path);
+    let mut options = std::fs::OpenOptions::new();
+    options.write(true).create_new(true);
+    #[cfg(unix)]
+    std::os::unix::fs::OpenOptionsExt::mode(&mut options, 0o600);
+    let mut file = options.open(&tmp_path)?;
+    file.write_all(contents)?;
+    file.sync_all()?;
+    file.set_permissions(permissions)?;
+    std::fs::rename(&tmp_path, path)
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
 
+    #[cfg(unix)]
+    #[test]
+    fn a_rewritten_config_keeps_its_permissions() {
+        use std::os::unix::fs::PermissionsExt;
+
+        let dir = tempfile::tempdir().unwrap();
+        let path = dir.path().join("ember.toml");
+        std::fs::write(&path, "requirepass = \"old\"\n").unwrap();
+        std::fs::set_permissions(&path, std::fs::Permissions::from_mode(0o600)).unwrap();
+
+        write_private(&path, b"requirepass = \"new\"\n").unwrap();
+        let mode = std::fs::metadata(&path).unwrap().permissions().mode();
+        assert_eq!(mode & 0o777, 0o600);
+        assert!(std::fs::read_to_string(&path).unwrap().contains("new"));
+    }
+
     #[test]
     fn parse_plain_bytes() {
         assert_eq!(parse_byte_size("1024").unwrap(), 1024);

crates/ember-server/src/connection/handler.rs

@@ -49,6 +49,7 @@ pub(super) async fn handle_frame_with_tx(
                     Ok(Command::Multi) => {
                         *tx_state = TransactionState::Queuing {
                             queue: Vec::new(),
+                            bytes: watched_keys.iter().map(|(key, _)| watch_size(key)).sum(),
                             error: false,
                         };
                         Frame::Simple("OK".into())
@@ -63,6 +64,11 @@ pub(super) async fn handle_frame_with_tx(
             } else if cmd_name == Some("WATCH") {
                 match Command::from_frame(frame) {
                     Ok(Command::Watch { keys }) => {
+                        let watched: usize = watched_keys.iter().map(|(k, _)| watch_size(k)).sum();
+                        let adding: usize = keys.iter().map(|k| watch_size(k)).sum();
+                        if watched + adding > ctx.limits.max_command_memory {
+                            return Frame::Error("ERR too many keys watched".into());
+                        }
                         // query current version for each key
                         for key in keys {
                             let ver = match engine
@@ -90,7 +96,11 @@ pub(super) async fn handle_frame_with_tx(
                 .await
             }
         }
-        TransactionState::Queuing { queue, error } => {
+        TransactionState::Queuing {
+            queue,
+            bytes,
+            error,
+        } => {
             match cmd_name {
                 Some("MULTI") => Frame::Error("ERR MULTI calls can not be nested".into()),
                 Some("WATCH") => Frame::Error("ERR WATCH inside MULTI is not allowed".into()),
@@ -149,6 +159,10 @@ pub(super) async fn handle_frame_with_tx(
                     queue.push(frame);
                     Frame::Simple("QUEUED".into())
                 }
+                _ if *bytes + request_size(&frame) > ctx.limits.max_command_memory => {
+                    *error = true;
+                    Frame::Error("ERR transaction too large; it will be discarded".into())
+                }
                 _ => {
                     // validate the command parses correctly before queuing
                     match Command::from_frame(frame.clone()) {
@@ -161,6 +175,7 @@ pub(super) async fn handle_frame_with_tx(
                                 )
                                 .await
                             } else {
+                                *bytes += request_size(&frame);
                                 queue.push(frame);
                                 Frame::Simple("QUEUED".into())
                             }
@@ -176,6 +191,26 @@ pub(super) async fn handle_frame_with_tx(
     }
 }
 
+/// The memory a watched key takes, counting its entry.
+fn watch_size(key: &str) -> usize {
+    key.len() + 32
+}
+
+/// The memory a queued request takes: its arguments, plus the overhead of
+/// each.
+fn request_size(frame: &Frame) -> usize {
+    match frame {
+        Frame::Array(args) => args
+            .iter()
+            .map(|arg| match arg {
+                Frame::Bulk(b) => b.len() + 32,
+                _ => 32,
+            })
+            .sum(),
+        _ => 32,
+    }
+}
+
 /// Re-queries all watched key versions and returns true if none have changed.
 async fn check_watched_keys(watched: &[(String, Option<u64>)], engine: &Engine) -> bool {
     for (key, original_ver) in watched {

crates/ember-server/src/connection/mod.rs

@@ -386,8 +386,12 @@ where
                         }
                     }
                 } else if session.is_authenticated() || is_allowed_before_auth(&frame) {
-                    let response = dispatch::process(
+                    // through the transaction handler, so MULTI, EXEC and
+                    // WATCH work in a batch with AUTH as they do without
+                    let response = handler::handle_frame_with_tx(
                         frame,
+                        &mut tx_state,
+                        &mut watched_keys,
                         &engine,
                         ctx,
                         slow_log,

crates/ember-server/src/connection_common.rs

@@ -479,8 +479,11 @@ pub enum TransactionState {
     /// Queuing mode — frames accumulate until EXEC or DISCARD.
     Queuing {
         queue: Vec<Frame>,
-        /// Set when a command in the queue had a parse error. EXEC returns
-        /// EXECABORT instead of executing.
+        /// Bytes the queue and the watched keys hold, which may not pass
+        /// `max_command_memory`: a transaction counts as one command.
+        bytes: usize,
+        /// Set when a command in the queue had a parse error, or the queue
+        /// grew too large. EXEC returns EXECABORT instead of executing.
         error: bool,
     },
 }

crates/ember-server/src/grpc/listener.rs

@@ -0,0 +1,136 @@
+//! The gRPC listener: accepts connections within `maxclients`, and runs
+//! the TLS handshake when TLS is configured.
+
+use std::io;
+use std::net::SocketAddr;
+use std::pin::Pin;
+use std::sync::Arc;
+use std::task::{Context, Poll};
+use std::time::Duration;
+
+use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
+use tokio::net::{TcpListener, TcpStream};
+use tokio::sync::{mpsc, OwnedSemaphorePermit, Semaphore};
+use tokio_rustls::server::TlsStream;
+use tokio_rustls::TlsAcceptor;
+use tokio_stream::wrappers::ReceiverStream;
+use tonic::transport::server::{Connected, TcpConnectInfo};
+use tracing::warn;
+
+/// An accepted gRPC connection. It holds one of the `maxclients` places
+/// while open, and reports its TCP addresses so protected mode can check
+/// the peer.
+pub(crate) struct Conn {
+    io: Io,
+    info: TcpConnectInfo,
+    _permit: OwnedSemaphorePermit,
+}
+
+enum Io {
+    Plain(TcpStream),
+    Tls(Box<TlsStream<TcpStream>>),
+}
+
+impl Connected for Conn {
+    type ConnectInfo = TcpConnectInfo;
+
+    fn connect_info(&self) -> TcpConnectInfo {
+        self.info.clone()
+    }
+}
+
+impl AsyncRead for Conn {
+    fn poll_read(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        buf: &mut ReadBuf<'_>,
+    ) -> Poll<io::Result<()>> {
+        match &mut self.get_mut().io {
+            Io::Plain(s) => Pin::new(s).poll_read(cx, buf),
+            Io::Tls(s) => Pin::new(s).poll_read(cx, buf),
+        }
+    }
+}
+
+impl AsyncWrite for Conn {
+    fn poll_write(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        buf: &[u8],
+    ) -> Poll<io::Result<usize>> {
+        match &mut self.get_mut().io {
+            Io::Plain(s) => Pin::new(s).poll_write(cx, buf),
+            Io::Tls(s) => Pin::new(s).poll_write(cx, buf),
+        }
+    }
+
+    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
+        match &mut self.get_mut().io {
+            Io::Plain(s) => Pin::new(s).poll_flush(cx),
+            Io::Tls(s) => Pin::new(s).poll_flush(cx),
+        }
+    }
+
+    fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
+        match &mut self.get_mut().io {
+            Io::Plain(s) => Pin::new(s).poll_shutdown(cx),
+            Io::Tls(s) => Pin::new(s).poll_shutdown(cx),
+        }
+    }
+}
+
+/// Accepts connections on `addr` for `Server::serve_with_incoming`.
+///
+/// Each connection takes a place from `clients`, the semaphore RESP
+/// connections share, and is closed when none is left. With `tls`, each
+/// handshake runs in its own task, so a slow client doesn't hold up the
+/// others.
+pub(crate) async fn incoming(
+    addr: SocketAddr,
+    tls: Option<TlsAcceptor>,
+    clients: Arc<Semaphore>,
+) -> io::Result<ReceiverStream<io::Result<Conn>>> {
+    let listener = TcpListener::bind(addr).await?;
+    let (tx, rx) = mpsc::channel(64);
+    tokio::spawn(async move {
+        while !tx.is_closed() {
+            let (stream, peer) = match listener.accept().await {
+                Ok(conn) => conn,
+                Err(e) => {
+                    warn!("gRPC accept error: {e}");
+                    tokio::time::sleep(Duration::from_millis(100)).await;
+                    continue;
+                }
+            };
+            let Ok(permit) = Arc::clone(&clients).try_acquire_owned() else {
+                warn!("connection limit reached, dropping gRPC connection from {peer}");
+                continue;
+            };
+            let _ = stream.set_nodelay(true);
+            let keepalive = socket2::TcpKeepalive::new().with_time(Duration::from_secs(60));
+            let _ = socket2::SockRef::from(&stream).set_tcp_keepalive(&keepalive);
+            let info = stream.connect_info();
+            let Some(acceptor) = tls.clone() else {
+                let conn = Conn {
+                    io: Io::Plain(stream),
+                    info,
+                    _permit: permit,
+                };
+                let _ = tx.send(Ok(conn)).await;
+                continue;
+            };
+            let tx = tx.clone();
+            tokio::spawn(async move {
+                if let Some(stream) = crate::server::tls_handshake(acceptor, stream, peer).await {
+                    let conn = Conn {
+                        io: Io::Tls(Box::new(stream)),
+                        info,
+                        _permit: permit,
+                    };
+                    let _ = tx.send(Ok(conn)).await;
+                }
+            });
+        }
+    });
+    Ok(ReceiverStream::new(rx))
+}

crates/ember-server/src/grpc/mod.rs

@@ -29,14 +29,14 @@ pub mod proto {
 
 mod hashes;
 mod keyspace;
+pub(crate) mod listener;
 mod lists;
 mod pubsub;
 mod reply;
 mod server;
 mod sets;
 mod sorted_sets;
 mod strings;
-pub(crate) mod tls;
 mod vector;
 
 use proto::ember_cache_server::EmberCache;