rustis 0.25.0

//! Minimal RESP3 responder, driven over any byte stream.
//!
//! It answers commands from a lookup table keyed on the command name, which is
//! enough to bring a real [`Client`](crate::client::Client) up: the handshake is
//! served like any other command, and a test only spells out the replies it
//! cares about. Unknown commands get an error reply, so a test that forgot one
//! fails on that error rather than hanging.
//!
//! Replies are written as raw bytes: the crate encodes requests only
//! ([`CommandEncoder`](crate::resp::CommandEncoder)), and writes a reply back
//! only from one it decoded ([`RawResponse`](crate::resp::RawResponse)), which
//! is a reply a test would have had to write in the first place.

use crate::{
    Future,
    client::{
        Config, CustomTransport, ServerConfig, TransportFactory, TransportReader, TransportWriter,
    },
};
use std::{
    collections::HashMap,
    sync::{
        Arc,
        atomic::{AtomicUsize, Ordering},
    },
};
use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};

/// Serves `server` on one end of a fresh in-memory pipe and hands the other end
/// back as a transport.
pub(crate) fn duplex_pair(server: FakeServer) -> (TransportReader, TransportWriter) {
    let (client_side, server_side) = tokio::io::duplex(4096);
    tokio::spawn(async move { server.serve(server_side).await });
    let (reader, writer) = tokio::io::split(client_side);
    (Box::new(reader), Box::new(writer))
}

/// A [`FakeServer`] listening on a real, ephemeral TCP port.
///
/// A pipe cannot stand in for a server a test must *address*: a cluster client
/// dials the host and port a `CLUSTER SHARDS` reply names, so the reply has to
/// name a port that answers. Dropping the node closes the listener, which is how
/// a test spells "this node is gone" — a later dial is refused rather than
/// hanging.
/// Binding and serving are two steps because a cluster reply has to name the
/// port it is served on: a test binds every node, builds the topology from the
/// ports it got, then starts answering with it.
pub(crate) struct TcpNode {
    pub addr: std::net::SocketAddr,
    accepted: Arc<AtomicUsize>,
    /// Held until [`serve`](Self::serve) hands it to the accept loop. Connections
    /// made in between wait in the backlog rather than being refused.
    listener: Option<tokio::net::TcpListener>,
    accept_loop: Option<tokio::task::JoinHandle<()>>,
}

impl TcpNode {
    /// Binds an ephemeral port, answering nothing yet.
    pub(crate) async fn bind() -> std::io::Result<Self> {
        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await?;
        Ok(Self {
            addr: listener.local_addr()?,
            accepted: Arc::new(AtomicUsize::new(0)),
            listener: Some(listener),
            accept_loop: None,
        })
    }

    /// Answers every connection from now on with `server`.
    pub(crate) fn serve(&mut self, server: FakeServer) {
        let Some(listener) = self.listener.take() else {
            panic!("a node serves once");
        };
        let accepted = Arc::clone(&self.accepted);

        self.accept_loop = Some(tokio::spawn(async move {
            while let Ok((stream, _)) = listener.accept().await {
                accepted.fetch_add(1, Ordering::SeqCst);
                let server = server.clone();
                tokio::spawn(async move { server.serve(stream).await });
            }
        }));
    }

    /// How many connections this node has accepted, which is what tells a test
    /// whether a node was dialled at all.
    pub(crate) fn accepted(&self) -> usize {
        self.accepted.load(Ordering::SeqCst)
    }

    /// The `(host, port)` pair a cluster configuration addresses this node by.
    pub(crate) fn address(&self) -> (String, u16) {
        (self.addr.ip().to_string(), self.addr.port())
    }
}

impl Drop for TcpNode {
    /// Closes the port, which is how a test spells "this node is gone": a later
    /// dial is refused rather than left hanging.
    fn drop(&mut self) {
        if let Some(accept_loop) = &self.accept_loop {
            accept_loop.abort();
        }
    }
}

/// A factory answering every dial with a fresh pipe, counting its calls so a
/// test can tell a reconnection from the initial connection.
struct DuplexTransport {
    server: FakeServer,
    dials: Arc<AtomicUsize>,
}

impl TransportFactory for DuplexTransport {
    fn connect(&self) -> Future<'_, (TransportReader, TransportWriter)> {
        let server = self.server.clone();
        let dials = Arc::clone(&self.dials);
        Box::pin(async move {
            dials.fetch_add(1, Ordering::SeqCst);
            Ok(duplex_pair(server))
        })
    }
}

pub(crate) fn duplex_config(server: FakeServer, dials: Arc<AtomicUsize>) -> Config {
    Config {
        server: ServerConfig::Custom(CustomTransport::new(DuplexTransport { server, dials })),
        ..Default::default()
    }
}

/// A RESP3 `HELLO` reply, the one every connection reads before anything else.
pub(crate) const HELLO_REPLY: &[u8] = b"%7\r\n\
$6\r\nserver\r\n$5\r\nredis\r\n\
$7\r\nversion\r\n$5\r\n7.4.0\r\n\
$5\r\nproto\r\n:3\r\n\
$2\r\nid\r\n:1\r\n\
$4\r\nmode\r\n$10\r\nstandalone\r\n\
$4\r\nrole\r\n$6\r\nmaster\r\n\
$7\r\nmodules\r\n*0\r\n";

/// A scripted set of replies, one per command name.
#[derive(Clone, Default)]
pub(crate) struct FakeServer {
    replies: HashMap<String, Vec<u8>>,
}

impl FakeServer {
    /// A server that answers the handshake and nothing else.
    pub(crate) fn new() -> Self {
        let mut replies = HashMap::new();
        replies.insert("HELLO".to_owned(), HELLO_REPLY.to_vec());
        Self { replies }
    }

    /// Answers `command` — matched case-insensitively, as Redis does — with the
    /// raw RESP bytes `reply`.
    pub(crate) fn reply(mut self, command: &str, reply: &[u8]) -> Self {
        self.replies
            .insert(command.to_ascii_uppercase(), reply.to_vec());
        self
    }

    /// Reads requests off `stream` and writes the scripted reply to each, until
    /// the peer closes the stream.
    pub(crate) async fn serve<S>(&self, mut stream: S)
    where
        S: AsyncRead + AsyncWrite + Unpin,
    {
        let mut buf = Vec::new();
        let mut chunk = [0u8; 1024];

        loop {
            let n = match stream.read(&mut chunk).await {
                Ok(0) | Err(_) => return,
                Ok(n) => n,
            };
            buf.extend_from_slice(&chunk[..n]);

            // A read can carry several pipelined requests, or half of one; only
            // whole requests are answered, and the remainder waits for the next
            // read.
            let mut cursor = 0;
            while let Some((name, next)) = parse_request(&buf[cursor..]) {
                cursor += next;
                let reply =
                    self.replies.get(&name).cloned().unwrap_or_else(|| {
                        format!("-ERR unknown command '{name}'\r\n").into_bytes()
                    });
                if stream.write_all(&reply).await.is_err() {
                    return;
                }
            }
            buf.drain(..cursor);
        }
    }
}

/// Parses one request — an array of bulk strings, which is the only shape a
/// client sends — returning its command name uppercased and its length in
/// bytes. `None` while the request is still incomplete.
fn parse_request(buf: &[u8]) -> Option<(String, usize)> {
    let (count, mut pos) = parse_prefixed_len(buf, b'*')?;

    let mut name = None;
    for _ in 0..count {
        let (len, next) = parse_prefixed_len(&buf[pos..], b'$')?;
        pos += next;
        let end = pos.checked_add(len)?;
        // The argument and its trailing CRLF must both be there.
        if buf.len() < end + 2 {
            return None;
        }
        if name.is_none() {
            name = Some(String::from_utf8_lossy(&buf[pos..end]).to_ascii_uppercase());
        }
        pos = end + 2;
    }

    Some((name?, pos))
}

/// Parses a `<tag><len>\r\n` header, returning the length and the header size.
fn parse_prefixed_len(buf: &[u8], tag: u8) -> Option<(usize, usize)> {
    if buf.first() != Some(&tag) {
        return None;
    }
    let crlf = buf.windows(2).position(|w| w == b"\r\n")?;
    let len = std::str::from_utf8(&buf[1..crlf]).ok()?.parse().ok()?;
    Some((len, crlf + 2))
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn a_request_is_parsed_into_its_command_name() {
        let (name, len) = parse_request(b"*2\r\n$3\r\nget\r\n$1\r\nk\r\n").unwrap();
        assert_eq!("GET", name);
        assert_eq!(20, len);
    }

    #[test]
    fn an_incomplete_request_is_not_parsed() {
        assert!(parse_request(b"*2\r\n$3\r\nget\r\n$1\r\n").is_none());
        assert!(parse_request(b"*2\r\n$3\r\nge").is_none());
        assert!(parse_request(b"*2\r\n").is_none());
    }

    #[tokio::test]
    async fn the_handshake_and_a_scripted_command_are_answered_in_order() {
        let (client_side, server_side) = tokio::io::duplex(1024);
        let server = FakeServer::new().reply("PING", b"+PONG\r\n");
        tokio::spawn(async move { server.serve(server_side).await });

        let mut client_side = client_side;
        client_side
            .write_all(b"*2\r\n$5\r\nhello\r\n$1\r\n3\r\n*1\r\n$4\r\nping\r\n")
            .await
            .unwrap();

        let mut out = vec![0u8; HELLO_REPLY.len() + 7];
        client_side.read_exact(&mut out).await.unwrap();
        assert_eq!(&out[..HELLO_REPLY.len()], HELLO_REPLY);
        assert_eq!(&out[HELLO_REPLY.len()..], b"+PONG\r\n");
    }
}