kevy_replicate/

handshake.rs

//! Replication handshake — the first round-trip a replica makes
//! against the primary's replication TCP listener.
//!
//! Wire shape:
//!
//! - Replica → primary: a RESP2 multi-bulk command
//!   `REPLICATE FROM <generation> <from-offset> ID <replica-id>`
//!   (6 args). `<generation>` is the feed generation the replica's
//!   DATA reflects (`0` = unknown / fresh); `<from-offset>` is `0`
//!   for a fresh replica (full-sync intent) or the last applied
//!   offset from a reconnecting replica. Offsets are only meaningful
//!   within one generation — the primary must never serve
//!   `<from-offset>` continuity across a generation mismatch (that
//!   is the offset-aliasing hole this field closes).
//! - Primary → replica: a RESP2 simple string
//!   `+ACK <generation> <current-offset>`, where `<generation>` is
//!   the primary's current feed generation and `<current-offset>`
//!   echoes the granted resume offset. The replica records both and
//!   starts consuming; if the primary's
//!   [`crate::source::ReplicationSource`] cannot serve from
//!   `<from-offset>` (TooOld), or the generations mismatch on a
//!   nonzero offset, the primary instead begins a snapshot ship
//!   (handled by the wiring layer, not here).
//!
//! This module owns only the parse + format primitives. Socket I/O,
//! retry, and "did the primary choose snapshot vs live stream" logic
//! live in the future replication source/replica modules.

use kevy_resp::Argv;

/// Parsed `REPLICATE FROM <generation> <from-offset> ID <replica-id>`
/// request.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct HandshakeReq {
    /// Feed generation the replica's data reflects. `0` = unknown /
    /// fresh — the primary treats it as "no continuity claim".
    pub generation: u64,
    /// Offset the replica wants to resume from. `0` = fresh replica.
    /// Only meaningful within `generation`.
    pub from_offset: u64,
    /// Replica-supplied identifier (operator-set, opaque to the
    /// primary other than for slot bookkeeping).
    pub replica_id: String,
}

/// Why a [`parse_replicate_from`] call rejected its input.
#[derive(Debug, PartialEq, Eq)]
pub enum HandshakeError {
    /// First arg is not "REPLICATE" (case-insensitive).
    BadCommand,
    /// Argument count is not exactly 6.
    WrongArity(usize),
    /// Second arg is not "FROM" (case-insensitive).
    BadFromKeyword,
    /// Third arg (generation) did not parse as an unsigned decimal `u64`.
    BadGeneration,
    /// Fourth arg (offset) did not parse as an unsigned decimal `u64`.
    BadOffset,
    /// Fifth arg is not "ID" (case-insensitive).
    BadIdKeyword,
    /// Replica id is empty or not valid UTF-8.
    BadReplicaId,
}

impl std::fmt::Display for HandshakeError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::BadCommand => write!(f, "expected REPLICATE command"),
            Self::WrongArity(n) => write!(f, "REPLICATE expects 6 args, got {n}"),
            Self::BadFromKeyword => write!(f, "expected 'FROM' keyword"),
            Self::BadGeneration => write!(f, "generation must be an unsigned decimal"),
            Self::BadOffset => write!(f, "from-offset must be an unsigned decimal"),
            Self::BadIdKeyword => write!(f, "expected 'ID' keyword"),
            Self::BadReplicaId => write!(f, "replica id must be non-empty UTF-8"),
        }
    }
}

impl std::error::Error for HandshakeError {}

/// Parse a `REPLICATE FROM <generation> <offset> ID <id>` command
/// from an already-decoded [`Argv`] (the caller has run the bytes
/// through `kevy_resp::parse_command_into` first).
// missing_panics_doc: the unwraps are guarded by the `len() != 6` arity check
// above them — unreachable, not a caller-facing panic condition.
#[allow(clippy::missing_panics_doc)]
pub fn parse_replicate_from(argv: &Argv) -> Result<HandshakeReq, HandshakeError> {
    if argv.len() != 6 {
        return Err(HandshakeError::WrongArity(argv.len()));
    }
    if !eq_ascii_ci(argv.get(0).unwrap(), b"REPLICATE") {
        return Err(HandshakeError::BadCommand);
    }
    if !eq_ascii_ci(argv.get(1).unwrap(), b"FROM") {
        return Err(HandshakeError::BadFromKeyword);
    }
    let generation =
        parse_decimal_u64(argv.get(2).unwrap()).ok_or(HandshakeError::BadGeneration)?;
    let from_offset = parse_decimal_u64(argv.get(3).unwrap()).ok_or(HandshakeError::BadOffset)?;
    if !eq_ascii_ci(argv.get(4).unwrap(), b"ID") {
        return Err(HandshakeError::BadIdKeyword);
    }
    let id_bytes = argv.get(5).unwrap();
    if id_bytes.is_empty() {
        return Err(HandshakeError::BadReplicaId);
    }
    let replica_id =
        std::str::from_utf8(id_bytes).map_err(|_| HandshakeError::BadReplicaId)?.to_string();
    Ok(HandshakeReq { generation, from_offset, replica_id })
}

/// Encode the primary's `+ACK <generation> <current-offset>\r\n`
/// response. `generation` is the primary's CURRENT feed generation —
/// the replica records it as the generation of whatever data this
/// session delivers (frames or snapshot).
pub fn encode_ack(generation: u64, current_offset: u64) -> Vec<u8> {
    let mut out = Vec::with_capacity(8 + 20 + 20 + 3);
    out.extend_from_slice(b"+ACK ");
    push_u64(&mut out, generation);
    out.push(b' ');
    push_u64(&mut out, current_offset);
    out.extend_from_slice(b"\r\n");
    out
}

fn eq_ascii_ci(a: &[u8], b: &[u8]) -> bool {
    a.len() == b.len() && a.iter().zip(b).all(|(x, y)| x.eq_ignore_ascii_case(y))
}

fn parse_decimal_u64(bytes: &[u8]) -> Option<u64> {
    if bytes.is_empty() {
        return None;
    }
    let mut n: u64 = 0;
    for &b in bytes {
        if !b.is_ascii_digit() {
            return None;
        }
        n = n.checked_mul(10)?.checked_add(u64::from(b - b'0'))?;
    }
    Some(n)
}

fn push_u64(out: &mut Vec<u8>, n: u64) {
    if n == 0 {
        out.push(b'0');
        return;
    }
    let mut tmp = [0u8; 20];
    let mut i = tmp.len();
    let mut v = n;
    while v != 0 {
        i -= 1;
        tmp[i] = b'0' + (v % 10) as u8;
        v /= 10;
    }
    out.extend_from_slice(&tmp[i..]);
}

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

    fn argv(args: &[&[u8]]) -> Argv {
        let mut a = Argv::default();
        for arg in args {
            a.push(arg);
        }
        a
    }

    #[test]
    fn parses_fresh_replica_from_zero() {
        let req =
            parse_replicate_from(&argv(&[b"REPLICATE", b"FROM", b"0", b"0", b"ID", b"replica-a"]))
                .unwrap();
        assert_eq!(req.generation, 0);
        assert_eq!(req.from_offset, 0);
        assert_eq!(req.replica_id, "replica-a");
    }

    #[test]
    fn parses_reconnect_with_generation_and_large_offset() {
        let req = parse_replicate_from(&argv(&[
            b"REPLICATE",
            b"FROM",
            b"7",
            b"4294967296", // 2^32 — guarantees u64 path
            b"ID",
            b"node-7",
        ]))
        .unwrap();
        assert_eq!(req.generation, 7);
        assert_eq!(req.from_offset, 4_294_967_296);
        assert_eq!(req.replica_id, "node-7");
    }

    #[test]
    fn keywords_are_case_insensitive() {
        let req =
            parse_replicate_from(&argv(&[b"replicate", b"from", b"2", b"1", b"id", b"x"])).unwrap();
        assert_eq!(req.generation, 2);
        assert_eq!(req.from_offset, 1);
        assert_eq!(req.replica_id, "x");
    }

    #[test]
    fn wrong_arity_rejected_with_actual_count() {
        // The legacy 5-arg (gen-less) form is a WrongArity rejection,
        // not a silent downgrade — 4.0 is a clean wire break.
        let err =
            parse_replicate_from(&argv(&[b"REPLICATE", b"FROM", b"0", b"ID", b"a"])).unwrap_err();
        assert_eq!(err, HandshakeError::WrongArity(5));
    }

    #[test]
    fn wrong_command_rejected() {
        let err = parse_replicate_from(&argv(&[b"SUBSCRIBE", b"FROM", b"0", b"0", b"ID", b"a"]))
            .unwrap_err();
        assert_eq!(err, HandshakeError::BadCommand);
    }

    #[test]
    fn wrong_from_keyword_rejected() {
        let err = parse_replicate_from(&argv(&[b"REPLICATE", b"AT", b"0", b"0", b"ID", b"a"]))
            .unwrap_err();
        assert_eq!(err, HandshakeError::BadFromKeyword);
    }

    #[test]
    fn wrong_id_keyword_rejected() {
        let err = parse_replicate_from(&argv(&[b"REPLICATE", b"FROM", b"0", b"0", b"NAME", b"a"]))
            .unwrap_err();
        assert_eq!(err, HandshakeError::BadIdKeyword);
    }

    #[test]
    fn non_decimal_generation_rejected() {
        let err = parse_replicate_from(&argv(&[b"REPLICATE", b"FROM", b"NaN", b"0", b"ID", b"a"]))
            .unwrap_err();
        assert_eq!(err, HandshakeError::BadGeneration);
    }

    #[test]
    fn non_decimal_offset_rejected() {
        let err = parse_replicate_from(&argv(&[b"REPLICATE", b"FROM", b"1", b"NaN", b"ID", b"a"]))
            .unwrap_err();
        assert_eq!(err, HandshakeError::BadOffset);
    }

    #[test]
    fn negative_offset_rejected_as_bad_offset() {
        let err = parse_replicate_from(&argv(&[b"REPLICATE", b"FROM", b"1", b"-1", b"ID", b"a"]))
            .unwrap_err();
        assert_eq!(err, HandshakeError::BadOffset);
    }

    #[test]
    fn empty_replica_id_rejected() {
        let err = parse_replicate_from(&argv(&[b"REPLICATE", b"FROM", b"0", b"0", b"ID", b""]))
            .unwrap_err();
        assert_eq!(err, HandshakeError::BadReplicaId);
    }

    #[test]
    fn non_utf8_replica_id_rejected() {
        let err = parse_replicate_from(&argv(&[
            b"REPLICATE",
            b"FROM",
            b"0",
            b"0",
            b"ID",
            &[0xFF, 0xFE, 0xFD], // invalid UTF-8
        ]))
        .unwrap_err();
        assert_eq!(err, HandshakeError::BadReplicaId);
    }

    #[test]
    fn ack_format_for_zero() {
        assert_eq!(encode_ack(1, 0), b"+ACK 1 0\r\n");
    }

    #[test]
    fn ack_format_for_large_values() {
        assert_eq!(encode_ack(12, 987_654_321), b"+ACK 12 987654321\r\n");
    }
}