1use kevy_resp::Argv;
22
23#[derive(Debug, Clone, PartialEq, Eq)]
25pub struct HandshakeReq {
26 pub from_offset: u64,
28 pub replica_id: String,
31}
32
33#[derive(Debug, PartialEq, Eq)]
35pub enum HandshakeError {
36 BadCommand,
38 WrongArity(usize),
40 BadFromKeyword,
42 BadOffset,
44 BadIdKeyword,
46 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#[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
98pub 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", 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], ]))
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}