Skip to main content

kevy_replicate/
handshake.rs

1//! Replication handshake — the first round-trip a replica makes
2//! against the primary's replication TCP listener.
3//!
4//! Wire shape:
5//!
6//! - Replica → primary: a RESP2 multi-bulk command
7//!   `REPLICATE FROM <from-offset> ID <replica-id>` (5 args). The
8//!   `<from-offset>` is `0` for a fresh replica (full-sync intent) or
9//!   the last applied offset from a reconnecting replica.
10//! - Primary → replica: a RESP2 simple string `+ACK <current-offset>`,
11//!   where `<current-offset>` is the primary's `next_offset` at the
12//!   moment of ack. The replica records it and starts consuming live
13//!   frames; if the primary's [`crate::source::ReplicationSource`]
14//!   cannot serve from `<from-offset>` (TooOld), the primary instead
15//!   begins a snapshot ship (handled by the wiring layer, not here).
16//!
17//! This module owns only the parse + format primitives. Socket I/O,
18//! retry, and "did the primary choose snapshot vs live stream" logic
19//! live in the future replication source/replica modules.
20
21use kevy_resp::Argv;
22
23/// Parsed `REPLICATE FROM <from-offset> ID <replica-id>` request.
24#[derive(Debug, Clone, PartialEq, Eq)]
25pub struct HandshakeReq {
26    /// Offset the replica wants to resume from. `0` = fresh replica.
27    pub from_offset: u64,
28    /// Replica-supplied identifier (operator-set, opaque to the
29    /// primary other than for slot bookkeeping).
30    pub replica_id: String,
31}
32
33/// Why a [`parse_replicate_from`] call rejected its input.
34#[derive(Debug, PartialEq, Eq)]
35pub enum HandshakeError {
36    /// First arg is not "REPLICATE" (case-insensitive).
37    BadCommand,
38    /// Argument count is not exactly 5.
39    WrongArity(usize),
40    /// Second arg is not "FROM" (case-insensitive).
41    BadFromKeyword,
42    /// Third arg did not parse as an unsigned decimal `u64`.
43    BadOffset,
44    /// Fourth arg is not "ID" (case-insensitive).
45    BadIdKeyword,
46    /// Replica id is empty or not valid UTF-8.
47    BadReplicaId,
48}
49
50impl std::fmt::Display for HandshakeError {
51    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
52        match self {
53            Self::BadCommand => write!(f, "expected REPLICATE command"),
54            Self::WrongArity(n) => write!(f, "REPLICATE expects 5 args, got {n}"),
55            Self::BadFromKeyword => write!(f, "expected 'FROM' keyword"),
56            Self::BadOffset => write!(f, "from-offset must be an unsigned decimal"),
57            Self::BadIdKeyword => write!(f, "expected 'ID' keyword"),
58            Self::BadReplicaId => write!(f, "replica id must be non-empty UTF-8"),
59        }
60    }
61}
62
63impl std::error::Error for HandshakeError {}
64
65/// Parse a `REPLICATE FROM <offset> ID <id>` command from an
66/// already-decoded [`Argv`] (the caller has run the bytes through
67/// `kevy_resp::parse_command_into` first).
68// missing_panics_doc: the unwraps are guarded by the `len() != 5` arity check
69// above them — unreachable, not a caller-facing panic condition.
70#[allow(clippy::missing_panics_doc)]
71pub fn parse_replicate_from(argv: &Argv) -> Result<HandshakeReq, HandshakeError> {
72    if argv.len() != 5 {
73        return Err(HandshakeError::WrongArity(argv.len()));
74    }
75    if !eq_ascii_ci(argv.get(0).unwrap(), b"REPLICATE") {
76        return Err(HandshakeError::BadCommand);
77    }
78    if !eq_ascii_ci(argv.get(1).unwrap(), b"FROM") {
79        return Err(HandshakeError::BadFromKeyword);
80    }
81    let from_offset =
82        parse_decimal_u64(argv.get(2).unwrap()).ok_or(HandshakeError::BadOffset)?;
83    if !eq_ascii_ci(argv.get(3).unwrap(), b"ID") {
84        return Err(HandshakeError::BadIdKeyword);
85    }
86    let id_bytes = argv.get(4).unwrap();
87    if id_bytes.is_empty() {
88        return Err(HandshakeError::BadReplicaId);
89    }
90    let replica_id =
91        std::str::from_utf8(id_bytes).map_err(|_| HandshakeError::BadReplicaId)?.to_string();
92    Ok(HandshakeReq {
93        from_offset,
94        replica_id,
95    })
96}
97
98/// Encode the primary's `+ACK <current-offset>\r\n` response.
99pub fn encode_ack(current_offset: u64) -> Vec<u8> {
100    let mut out = Vec::with_capacity(8 + 20 + 2);
101    out.extend_from_slice(b"+ACK ");
102    push_u64(&mut out, current_offset);
103    out.extend_from_slice(b"\r\n");
104    out
105}
106
107fn eq_ascii_ci(a: &[u8], b: &[u8]) -> bool {
108    a.len() == b.len()
109        && a.iter()
110            .zip(b)
111            .all(|(x, y)| x.eq_ignore_ascii_case(y))
112}
113
114fn parse_decimal_u64(bytes: &[u8]) -> Option<u64> {
115    if bytes.is_empty() {
116        return None;
117    }
118    let mut n: u64 = 0;
119    for &b in bytes {
120        if !b.is_ascii_digit() {
121            return None;
122        }
123        n = n.checked_mul(10)?.checked_add(u64::from(b - b'0'))?;
124    }
125    Some(n)
126}
127
128fn push_u64(out: &mut Vec<u8>, n: u64) {
129    if n == 0 {
130        out.push(b'0');
131        return;
132    }
133    let mut tmp = [0u8; 20];
134    let mut i = tmp.len();
135    let mut v = n;
136    while v != 0 {
137        i -= 1;
138        tmp[i] = b'0' + (v % 10) as u8;
139        v /= 10;
140    }
141    out.extend_from_slice(&tmp[i..]);
142}
143
144#[cfg(test)]
145mod tests {
146    use super::*;
147
148    fn argv(args: &[&[u8]]) -> Argv {
149        let mut a = Argv::default();
150        for arg in args {
151            a.push(arg);
152        }
153        a
154    }
155
156    #[test]
157    fn parses_fresh_replica_from_zero() {
158        let req = parse_replicate_from(&argv(&[
159            b"REPLICATE",
160            b"FROM",
161            b"0",
162            b"ID",
163            b"replica-a",
164        ]))
165        .unwrap();
166        assert_eq!(req.from_offset, 0);
167        assert_eq!(req.replica_id, "replica-a");
168    }
169
170    #[test]
171    fn parses_reconnect_with_large_offset() {
172        let req = parse_replicate_from(&argv(&[
173            b"REPLICATE",
174            b"FROM",
175            b"4294967296", // 2^32 — guarantees u64 path
176            b"ID",
177            b"node-7",
178        ]))
179        .unwrap();
180        assert_eq!(req.from_offset, 4_294_967_296);
181        assert_eq!(req.replica_id, "node-7");
182    }
183
184    #[test]
185    fn keywords_are_case_insensitive() {
186        let req = parse_replicate_from(&argv(&[
187            b"replicate", b"from", b"1", b"id", b"x",
188        ]))
189        .unwrap();
190        assert_eq!(req.from_offset, 1);
191        assert_eq!(req.replica_id, "x");
192    }
193
194    #[test]
195    fn wrong_arity_rejected_with_actual_count() {
196        let err = parse_replicate_from(&argv(&[b"REPLICATE", b"FROM", b"0"])).unwrap_err();
197        assert_eq!(err, HandshakeError::WrongArity(3));
198    }
199
200    #[test]
201    fn wrong_command_rejected() {
202        let err =
203            parse_replicate_from(&argv(&[b"SUBSCRIBE", b"FROM", b"0", b"ID", b"a"])).unwrap_err();
204        assert_eq!(err, HandshakeError::BadCommand);
205    }
206
207    #[test]
208    fn wrong_from_keyword_rejected() {
209        let err =
210            parse_replicate_from(&argv(&[b"REPLICATE", b"AT", b"0", b"ID", b"a"])).unwrap_err();
211        assert_eq!(err, HandshakeError::BadFromKeyword);
212    }
213
214    #[test]
215    fn wrong_id_keyword_rejected() {
216        let err = parse_replicate_from(&argv(&[
217            b"REPLICATE",
218            b"FROM",
219            b"0",
220            b"NAME",
221            b"a",
222        ]))
223        .unwrap_err();
224        assert_eq!(err, HandshakeError::BadIdKeyword);
225    }
226
227    #[test]
228    fn non_decimal_offset_rejected() {
229        let err =
230            parse_replicate_from(&argv(&[b"REPLICATE", b"FROM", b"NaN", b"ID", b"a"]))
231                .unwrap_err();
232        assert_eq!(err, HandshakeError::BadOffset);
233    }
234
235    #[test]
236    fn negative_offset_rejected_as_bad_offset() {
237        let err =
238            parse_replicate_from(&argv(&[b"REPLICATE", b"FROM", b"-1", b"ID", b"a"]))
239                .unwrap_err();
240        assert_eq!(err, HandshakeError::BadOffset);
241    }
242
243    #[test]
244    fn empty_replica_id_rejected() {
245        let err =
246            parse_replicate_from(&argv(&[b"REPLICATE", b"FROM", b"0", b"ID", b""]))
247                .unwrap_err();
248        assert_eq!(err, HandshakeError::BadReplicaId);
249    }
250
251    #[test]
252    fn non_utf8_replica_id_rejected() {
253        let err = parse_replicate_from(&argv(&[
254            b"REPLICATE",
255            b"FROM",
256            b"0",
257            b"ID",
258            &[0xFF, 0xFE, 0xFD], // invalid UTF-8
259        ]))
260        .unwrap_err();
261        assert_eq!(err, HandshakeError::BadReplicaId);
262    }
263
264    #[test]
265    fn ack_format_for_zero() {
266        assert_eq!(encode_ack(0), b"+ACK 0\r\n");
267    }
268
269    #[test]
270    fn ack_format_for_large_offset() {
271        assert_eq!(encode_ack(987_654_321), b"+ACK 987654321\r\n");
272    }
273}