dynomite/proto/memcache/
parser.rs

1//! Memcached text-protocol parser.
2//!
3//! The parser is a single byte-driven state machine for requests
4//! and another for responses. It mutates a [`Msg`] in place: the
5//! command tag, key list, value-length accumulator, and parser
6//! cursor are written back so the streaming caller can resume the
7//! machine on more bytes without re-allocating state.
8//!
9//! The parsers MUST NOT panic on any input. Invalid bytes are
10//! reported via [`MsgParseResult::Error`].
11
12// The parser truncates ASCII-decimal accumulators into fixed-width
13// counters (`u32` for vlen, `usize` for cursor offsets). The
14// allowance covers these intentional `as` casts.
15#![allow(clippy::cast_possible_truncation)]
16#![allow(clippy::too_many_arguments)]
17#![allow(clippy::match_same_arms)]
18#![allow(clippy::needless_continue)]
19// The state machine deliberately keeps the `if (token == NULL)`
20// guard pattern; rewriting as `let-else` would collapse two
21// branches that are treated independently.
22#![allow(clippy::manual_let_else)]
23#![allow(clippy::redundant_else)]
24
25use super::commands::{
26    memcache_arithmetic, memcache_cas, memcache_delete, memcache_retrieval, memcache_storage,
27    memcache_touch,
28};
29use crate::msg::{KeyPos, Msg, MsgParseResult, MsgType};
30
31/// Maximum allowed Memcached key length in bytes.
32pub const MEMCACHE_MAX_KEY_LENGTH: usize = 250;
33
34/// Optional hash tag delimiters. When set, parsed keys carry the
35/// inner range between the delimiters as the routing tag.
36#[derive(Copy, Clone, Debug, Default, Eq, PartialEq)]
37pub struct HashTag {
38    /// Opening byte of the hash tag.
39    pub open: u8,
40    /// Closing byte of the hash tag.
41    pub close: u8,
42}
43
44/// State alphabet for [`memcache_parse_req`].
45///
46/// The numeric values are stable request state indices so external
47/// parity tooling can compare them directly.
48#[derive(Copy, Clone, Debug, Default, Eq, PartialEq)]
49#[repr(u32)]
50pub enum ReqState {
51    /// Initial state.
52    #[default]
53    Start = 0,
54    /// Reading the command keyword.
55    ReqType = 1,
56    /// Skipping spaces before the first key.
57    SpacesBeforeKey = 2,
58    /// Reading a key.
59    Key = 3,
60    /// Skipping spaces between keys (`get`/`gets`).
61    SpacesBeforeKeys = 4,
62    /// Skipping spaces before the storage flags field.
63    SpacesBeforeFlags = 5,
64    /// Reading the storage flags field.
65    Flags = 6,
66    /// Skipping spaces before the storage expiry field.
67    SpacesBeforeExpiry = 7,
68    /// Reading the storage expiry field.
69    Expiry = 8,
70    /// Skipping spaces before the value-length field.
71    SpacesBeforeVlen = 9,
72    /// Reading the value-length field.
73    Vlen = 10,
74    /// Skipping spaces before the CAS unique field.
75    SpacesBeforeCas = 11,
76    /// Reading the CAS unique field.
77    Cas = 12,
78    /// Awaiting LF before the value bytes.
79    RuntoVal = 13,
80    /// Consuming the value bytes.
81    Val = 14,
82    /// Skipping spaces before the arithmetic numeric argument.
83    SpacesBeforeNum = 15,
84    /// Reading the arithmetic numeric argument.
85    Num = 16,
86    /// Eating optional trailing bytes up to CR.
87    RuntoCrlf = 17,
88    /// Awaiting trailing CR.
89    Crlf = 18,
90    /// Reading the optional `noreply` token.
91    Noreply = 19,
92    /// State after consuming `noreply`.
93    AfterNoreply = 20,
94    /// Awaiting the trailing LF that terminates the request.
95    AlmostDone = 21,
96}
97
98impl ReqState {
99    fn from_u32(v: u32) -> Self {
100        match v {
101            1 => Self::ReqType,
102            2 => Self::SpacesBeforeKey,
103            3 => Self::Key,
104            4 => Self::SpacesBeforeKeys,
105            5 => Self::SpacesBeforeFlags,
106            6 => Self::Flags,
107            7 => Self::SpacesBeforeExpiry,
108            8 => Self::Expiry,
109            9 => Self::SpacesBeforeVlen,
110            10 => Self::Vlen,
111            11 => Self::SpacesBeforeCas,
112            12 => Self::Cas,
113            13 => Self::RuntoVal,
114            14 => Self::Val,
115            15 => Self::SpacesBeforeNum,
116            16 => Self::Num,
117            17 => Self::RuntoCrlf,
118            18 => Self::Crlf,
119            19 => Self::Noreply,
120            20 => Self::AfterNoreply,
121            21 => Self::AlmostDone,
122            _ => Self::Start,
123        }
124    }
125}
126
127/// State alphabet for [`memcache_parse_rsp`].
128#[derive(Copy, Clone, Debug, Default, Eq, PartialEq)]
129#[repr(u32)]
130pub enum RspState {
131    /// Initial state.
132    #[default]
133    Start = 0,
134    /// Reading a numeric response (`incr`/`decr` reply).
135    RspNum = 1,
136    /// Reading a textual response keyword.
137    RspStr = 2,
138    /// Skipping spaces before the key in a `VALUE` reply.
139    SpacesBeforeKey = 3,
140    /// Reading the key portion of a `VALUE` reply.
141    Key = 4,
142    /// Skipping spaces before the flags field.
143    SpacesBeforeFlags = 5,
144    /// Reading the flags field.
145    Flags = 6,
146    /// Skipping spaces before the value-length field.
147    SpacesBeforeVlen = 7,
148    /// Reading the value-length field.
149    Vlen = 8,
150    /// Awaiting LF before the value bytes.
151    RuntoVal = 9,
152    /// Consuming the value bytes.
153    Val = 10,
154    /// Awaiting LF after the value.
155    ValLf = 11,
156    /// Reading the trailing `END` token.
157    End = 12,
158    /// Eating optional trailing bytes up to CR.
159    RuntoCrlf = 13,
160    /// Awaiting trailing CR.
161    Crlf = 14,
162    /// Awaiting the trailing LF that terminates the response.
163    AlmostDone = 15,
164}
165
166impl RspState {
167    fn from_u32(v: u32) -> Self {
168        match v {
169            1 => Self::RspNum,
170            2 => Self::RspStr,
171            3 => Self::SpacesBeforeKey,
172            4 => Self::Key,
173            5 => Self::SpacesBeforeFlags,
174            6 => Self::Flags,
175            7 => Self::SpacesBeforeVlen,
176            8 => Self::Vlen,
177            9 => Self::RuntoVal,
178            10 => Self::Val,
179            11 => Self::ValLf,
180            12 => Self::End,
181            13 => Self::RuntoCrlf,
182            14 => Self::Crlf,
183            15 => Self::AlmostDone,
184            _ => Self::Start,
185        }
186    }
187}
188
189const CR: u8 = b'\r';
190const LF: u8 = b'\n';
191
192fn classify_command(token: &[u8]) -> MsgType {
193    match token {
194        b"get" => MsgType::ReqMcGet,
195        b"gets" => MsgType::ReqMcGets,
196        b"set" => MsgType::ReqMcSet,
197        b"add" => MsgType::ReqMcAdd,
198        b"cas" => MsgType::ReqMcCas,
199        b"incr" => MsgType::ReqMcIncr,
200        b"decr" => MsgType::ReqMcDecr,
201        b"quit" => MsgType::ReqMcQuit,
202        b"touch" => MsgType::ReqMcTouch,
203        b"append" => MsgType::ReqMcAppend,
204        b"delete" => MsgType::ReqMcDelete,
205        b"prepend" => MsgType::ReqMcPrepend,
206        b"replace" => MsgType::ReqMcReplace,
207        _ => MsgType::Unknown,
208    }
209}
210
211fn classify_response(token: &[u8]) -> MsgType {
212    match token {
213        b"END" => MsgType::RspMcEnd,
214        b"VALUE" => MsgType::RspMcValue,
215        b"ERROR" => MsgType::RspMcError,
216        b"STORED" => MsgType::RspMcStored,
217        b"EXISTS" => MsgType::RspMcExists,
218        b"DELETED" => MsgType::RspMcDeleted,
219        b"TOUCHED" => MsgType::RspMcTouched,
220        b"NOT_FOUND" => MsgType::RspMcNotFound,
221        b"NOT_STORED" => MsgType::RspMcNotStored,
222        b"CLIENT_ERROR" => MsgType::RspMcClientError,
223        b"SERVER_ERROR" => MsgType::RspMcServerError,
224        _ => MsgType::Unknown,
225    }
226}
227
228fn make_keypos(input: &[u8], start: usize, end: usize, hash_tag: Option<HashTag>) -> KeyPos {
229    let bytes = input[start..end].to_vec();
230    if let Some(tag) = hash_tag {
231        if let Some(open_idx) = bytes.iter().position(|&b| b == tag.open) {
232            if let Some(close_offset) = bytes[open_idx + 1..].iter().position(|&b| b == tag.close) {
233                let tag_start = open_idx + 1;
234                let tag_end = open_idx + 1 + close_offset;
235                return KeyPos::new(bytes, tag_start..tag_end);
236            }
237        }
238    }
239    KeyPos::without_tag(bytes)
240}
241
242/// Parse a Memcached request from `input` and update `r` in place.
243///
244/// On success, `r.ty()` is set
245/// to the recognised command, parsed keys are appended to
246/// [`Msg::keys`], and the parser cursor (`parser_pos`) advances
247/// just past the trailing LF. On truncated input the function
248/// returns [`MsgParseResult::Again`] and stores the partial state
249/// on `r` for resumption. Invalid bytes return
250/// [`MsgParseResult::Error`].
251///
252/// `hash_tag`, when set, configures the routing-tag delimiters used
253/// when populating [`KeyPos::tag`].
254///
255/// # Examples
256///
257/// ```
258/// use dynomite::msg::{Msg, MsgParseResult, MsgType};
259/// use dynomite::proto::memcache::memcache_parse_req;
260///
261/// let mut r = Msg::new(0, MsgType::Unknown, true);
262/// let res = memcache_parse_req(&mut r, b"set foo 0 0 3\r\nbar\r\n");
263/// assert_eq!(res, MsgParseResult::Ok);
264/// assert_eq!(r.ty(), MsgType::ReqMcSet);
265/// assert_eq!(r.keys()[0].key(), b"foo");
266/// assert_eq!(r.vlen(), 3);
267/// ```
268///
269/// The state machine intentionally lives in a single function so
270/// the byte-by-byte control flow stays in one place.
271#[allow(clippy::too_many_lines)]
272pub fn memcache_parse_req(r: &mut Msg, input: &[u8]) -> MsgParseResult {
273    memcache_parse_req_tagged(r, input, None)
274}
275
276/// Variant of [`memcache_parse_req`] that accepts an explicit
277/// hash-tag configuration.
278///
279/// # Examples
280///
281/// ```
282/// use dynomite::msg::{Msg, MsgType};
283/// use dynomite::proto::memcache::parser::{memcache_parse_req_tagged, HashTag};
284///
285/// let mut r = Msg::new(0, MsgType::Unknown, true);
286/// let tag = Some(HashTag { open: b'{', close: b'}' });
287/// let _ = memcache_parse_req_tagged(&mut r, b"get {abc}xyz\r\n", tag);
288/// assert_eq!(r.keys()[0].tag_bytes(), b"abc");
289/// ```
290#[allow(clippy::too_many_lines)]
291pub fn memcache_parse_req_tagged(
292    r: &mut Msg,
293    input: &[u8],
294    hash_tag: Option<HashTag>,
295) -> MsgParseResult {
296    if !r.is_request() {
297        r.set_parse_result(MsgParseResult::Error);
298        return MsgParseResult::Error;
299    }
300    let mut state = ReqState::from_u32(r.parser_state());
301    let mut p = r.parser_pos();
302    let mut token: Option<usize> = r.parser_token();
303    let mut vlen = r.vlen();
304    let mut ty = r.ty();
305    let mut is_read = r.flags().is_read;
306    let mut quit = r.flags().quit;
307    let mut expect_reply = r.flags().expect_datastore_reply;
308    let mut ntokens = r.ntokens();
309
310    'machine: while p < input.len() {
311        let ch = input[p];
312        match state {
313            ReqState::Start => {
314                if ch == b' ' {
315                    p += 1;
316                    continue;
317                }
318                if !ch.is_ascii_lowercase() {
319                    return finish_error(r, state, p, token, vlen, ty, is_read, quit, ntokens);
320                }
321                token = Some(p);
322                state = ReqState::ReqType;
323                // Do not advance; re-enter ReqType on the same byte.
324            }
325            ReqState::ReqType => {
326                if ch == b' ' || ch == CR {
327                    let start = match token {
328                        Some(s) => s,
329                        None => {
330                            return finish_error(
331                                r, state, p, token, vlen, ty, is_read, quit, ntokens,
332                            );
333                        }
334                    };
335                    let cmd = &input[start..p];
336                    token = None;
337                    ty = classify_command(cmd);
338                    ntokens = ntokens.saturating_add(1);
339                    is_read = matches!(
340                        ty,
341                        MsgType::ReqMcGet | MsgType::ReqMcGets | MsgType::ReqMcQuit
342                    );
343                    if matches!(ty, MsgType::ReqMcQuit) {
344                        quit = true;
345                        // Set state to SW_CRLF and step p back by one.
346                        state = ReqState::Crlf;
347                        // Do not advance; re-enter on this same byte.
348                        continue;
349                    }
350                    match ty {
351                        MsgType::ReqMcGet
352                        | MsgType::ReqMcGets
353                        | MsgType::ReqMcDelete
354                        | MsgType::ReqMcCas
355                        | MsgType::ReqMcSet
356                        | MsgType::ReqMcAdd
357                        | MsgType::ReqMcReplace
358                        | MsgType::ReqMcAppend
359                        | MsgType::ReqMcPrepend
360                        | MsgType::ReqMcIncr
361                        | MsgType::ReqMcDecr
362                        | MsgType::ReqMcTouch => {
363                            if ch == CR {
364                                return finish_error(
365                                    r, state, p, token, vlen, ty, is_read, quit, ntokens,
366                                );
367                            }
368                            state = ReqState::SpacesBeforeKey;
369                            p += 1;
370                            continue;
371                        }
372                        MsgType::Unknown => {
373                            return finish_error(
374                                r, state, p, token, vlen, ty, is_read, quit, ntokens,
375                            );
376                        }
377                        _ => {
378                            return finish_error(
379                                r, state, p, token, vlen, ty, is_read, quit, ntokens,
380                            );
381                        }
382                    }
383                } else if !ch.is_ascii_lowercase() {
384                    return finish_error(r, state, p, token, vlen, ty, is_read, quit, ntokens);
385                } else {
386                    p += 1;
387                }
388            }
389            ReqState::SpacesBeforeKey => {
390                if ch == b' ' {
391                    p += 1;
392                } else {
393                    token = None;
394                    state = ReqState::Key;
395                    // Do not advance; re-process this byte under Key state.
396                }
397            }
398            ReqState::Key => {
399                if token.is_none() {
400                    token = Some(p);
401                }
402                if ch == b' ' || ch == CR {
403                    let start = token.expect("token recorded");
404                    let keylen = p - start;
405                    if keylen == 0 || keylen > MEMCACHE_MAX_KEY_LENGTH {
406                        return finish_error(r, state, p, token, vlen, ty, is_read, quit, ntokens);
407                    }
408                    let kp = make_keypos(input, start, p, hash_tag);
409                    r.push_key(kp);
410                    ntokens = ntokens.saturating_add(1);
411                    token = None;
412                    let storage = memcache_storage(ty);
413                    let arithmetic = memcache_arithmetic(ty);
414                    let touch = memcache_touch(ty);
415                    let delete = memcache_delete(ty);
416                    let retrieval = memcache_retrieval(ty);
417                    if storage {
418                        state = ReqState::SpacesBeforeFlags;
419                    } else if arithmetic || touch {
420                        state = ReqState::SpacesBeforeNum;
421                    } else if delete {
422                        state = ReqState::RuntoCrlf;
423                    } else if retrieval {
424                        state = ReqState::SpacesBeforeKeys;
425                    } else {
426                        state = ReqState::RuntoCrlf;
427                    }
428                    if ch == CR {
429                        if storage || arithmetic {
430                            return finish_error(
431                                r, state, p, token, vlen, ty, is_read, quit, ntokens,
432                            );
433                        }
434                        // Re-enter on the CR byte (do not advance).
435                    } else {
436                        p += 1;
437                    }
438                } else {
439                    p += 1;
440                }
441            }
442            ReqState::SpacesBeforeKeys => {
443                match ch {
444                    b' ' => {
445                        p += 1;
446                    }
447                    CR => {
448                        state = ReqState::AlmostDone;
449                        p += 1;
450                    }
451                    _ => {
452                        token = None;
453                        state = ReqState::Key;
454                        // Do not advance.
455                    }
456                }
457            }
458            ReqState::SpacesBeforeFlags => {
459                if ch == b' ' {
460                    p += 1;
461                } else if ch.is_ascii_digit() {
462                    token = Some(p);
463                    state = ReqState::Flags;
464                    p += 1;
465                } else {
466                    return finish_error(r, state, p, token, vlen, ty, is_read, quit, ntokens);
467                }
468            }
469            ReqState::Flags => {
470                if ch.is_ascii_digit() {
471                    p += 1;
472                } else if ch == b' ' {
473                    token = None;
474                    state = ReqState::SpacesBeforeExpiry;
475                    p += 1;
476                } else {
477                    return finish_error(r, state, p, token, vlen, ty, is_read, quit, ntokens);
478                }
479            }
480            ReqState::SpacesBeforeExpiry => {
481                if ch == b' ' {
482                    p += 1;
483                } else if ch.is_ascii_digit() {
484                    token = Some(p);
485                    state = ReqState::Expiry;
486                    p += 1;
487                } else {
488                    return finish_error(r, state, p, token, vlen, ty, is_read, quit, ntokens);
489                }
490            }
491            ReqState::Expiry => {
492                if ch.is_ascii_digit() {
493                    p += 1;
494                } else if ch == b' ' {
495                    token = None;
496                    state = ReqState::SpacesBeforeVlen;
497                    p += 1;
498                } else {
499                    return finish_error(r, state, p, token, vlen, ty, is_read, quit, ntokens);
500                }
501            }
502            ReqState::SpacesBeforeVlen => {
503                if ch == b' ' {
504                    p += 1;
505                } else if ch.is_ascii_digit() {
506                    vlen = u32::from(ch - b'0');
507                    state = ReqState::Vlen;
508                    p += 1;
509                } else {
510                    return finish_error(r, state, p, token, vlen, ty, is_read, quit, ntokens);
511                }
512            }
513            ReqState::Vlen => {
514                if ch.is_ascii_digit() {
515                    vlen = vlen.saturating_mul(10).saturating_add(u32::from(ch - b'0'));
516                    p += 1;
517                } else if memcache_cas(ty) {
518                    if ch != b' ' {
519                        return finish_error(r, state, p, token, vlen, ty, is_read, quit, ntokens);
520                    }
521                    token = None;
522                    state = ReqState::SpacesBeforeCas;
523                    // Do not advance; re-enter on the same byte.
524                } else if ch == b' ' || ch == CR {
525                    token = None;
526                    state = ReqState::RuntoCrlf;
527                    // Do not advance.
528                } else {
529                    return finish_error(r, state, p, token, vlen, ty, is_read, quit, ntokens);
530                }
531            }
532            ReqState::SpacesBeforeCas => {
533                if ch == b' ' {
534                    p += 1;
535                } else if ch.is_ascii_digit() {
536                    token = Some(p);
537                    state = ReqState::Cas;
538                    p += 1;
539                } else {
540                    return finish_error(r, state, p, token, vlen, ty, is_read, quit, ntokens);
541                }
542            }
543            ReqState::Cas => {
544                if ch.is_ascii_digit() {
545                    p += 1;
546                } else if ch == b' ' || ch == CR {
547                    token = None;
548                    state = ReqState::RuntoCrlf;
549                    // Do not advance.
550                } else {
551                    return finish_error(r, state, p, token, vlen, ty, is_read, quit, ntokens);
552                }
553            }
554            ReqState::RuntoVal => match ch {
555                LF => {
556                    state = ReqState::Val;
557                    p += 1;
558                }
559                _ => {
560                    return finish_error(r, state, p, token, vlen, ty, is_read, quit, ntokens);
561                }
562            },
563            ReqState::Val => {
564                let m = p.saturating_add(vlen as usize);
565                if m >= input.len() {
566                    let consumed = input.len() - p;
567                    vlen = vlen.saturating_sub(consumed as u32);
568                    p = input.len();
569                    break 'machine;
570                }
571                if input[m] != CR {
572                    return finish_error(r, state, p, token, vlen, ty, is_read, quit, ntokens);
573                }
574                p = m + 1;
575                state = ReqState::AlmostDone;
576            }
577            ReqState::SpacesBeforeNum => {
578                if ch == b' ' {
579                    p += 1;
580                } else if ch.is_ascii_digit() || ch == b'-' {
581                    token = Some(p);
582                    state = ReqState::Num;
583                    p += 1;
584                } else {
585                    return finish_error(r, state, p, token, vlen, ty, is_read, quit, ntokens);
586                }
587            }
588            ReqState::Num => {
589                if ch.is_ascii_digit() {
590                    p += 1;
591                } else if ch == b' ' || ch == CR {
592                    token = None;
593                    state = ReqState::RuntoCrlf;
594                    // Do not advance.
595                } else {
596                    return finish_error(r, state, p, token, vlen, ty, is_read, quit, ntokens);
597                }
598            }
599            ReqState::RuntoCrlf => match ch {
600                b' ' => {
601                    p += 1;
602                }
603                b'n' if (memcache_storage(ty)
604                    || memcache_arithmetic(ty)
605                    || memcache_delete(ty)
606                    || memcache_touch(ty)) =>
607                {
608                    token = Some(p);
609                    state = ReqState::Noreply;
610                    p += 1;
611                }
612                CR => {
613                    if memcache_storage(ty) {
614                        state = ReqState::RuntoVal;
615                    } else {
616                        state = ReqState::AlmostDone;
617                    }
618                    p += 1;
619                }
620                _ => {
621                    return finish_error(r, state, p, token, vlen, ty, is_read, quit, ntokens);
622                }
623            },
624            ReqState::Noreply => match ch {
625                b' ' | CR => {
626                    let start = match token {
627                        Some(s) => s,
628                        None => {
629                            return finish_error(
630                                r, state, p, token, vlen, ty, is_read, quit, ntokens,
631                            );
632                        }
633                    };
634                    if p - start == 7 && &input[start..p] == b"noreply" {
635                        token = None;
636                        expect_reply = false;
637                        state = ReqState::AfterNoreply;
638                        // Do not advance.
639                    } else {
640                        return finish_error(r, state, p, token, vlen, ty, is_read, quit, ntokens);
641                    }
642                }
643                _ => {
644                    p += 1;
645                }
646            },
647            ReqState::AfterNoreply => match ch {
648                b' ' => {
649                    p += 1;
650                }
651                CR => {
652                    if memcache_storage(ty) {
653                        state = ReqState::RuntoVal;
654                    } else {
655                        state = ReqState::AlmostDone;
656                    }
657                    p += 1;
658                }
659                _ => {
660                    return finish_error(r, state, p, token, vlen, ty, is_read, quit, ntokens);
661                }
662            },
663            ReqState::Crlf => match ch {
664                b' ' => {
665                    p += 1;
666                }
667                CR => {
668                    state = ReqState::AlmostDone;
669                    p += 1;
670                }
671                _ => {
672                    return finish_error(r, state, p, token, vlen, ty, is_read, quit, ntokens);
673                }
674            },
675            ReqState::AlmostDone => match ch {
676                LF => {
677                    return finish_done(r, p + 1, ty, is_read, quit, expect_reply, ntokens, vlen);
678                }
679                _ => {
680                    return finish_error(r, state, p, token, vlen, ty, is_read, quit, ntokens);
681                }
682            },
683        }
684    }
685
686    // Reached end of input without completing.
687    r.set_parser_state(state as u32);
688    r.set_parser_pos(p);
689    r.set_parser_token(token);
690    r.set_vlen(vlen);
691    r.set_ntokens(ntokens);
692    if ty != MsgType::Unknown {
693        r.set_type(ty);
694    }
695    r.flags_mut().is_read = is_read;
696    r.flags_mut().quit = quit;
697    r.flags_mut().expect_datastore_reply = expect_reply;
698    r.set_parse_result(MsgParseResult::Again);
699    MsgParseResult::Again
700}
701
702#[allow(clippy::too_many_arguments)]
703fn finish_done(
704    r: &mut Msg,
705    next_pos: usize,
706    ty: MsgType,
707    is_read: bool,
708    quit: bool,
709    expect_reply: bool,
710    ntokens: u32,
711    vlen: u32,
712) -> MsgParseResult {
713    r.set_type(ty);
714    r.flags_mut().is_read = is_read;
715    r.flags_mut().quit = quit;
716    r.flags_mut().expect_datastore_reply = expect_reply;
717    r.set_ntokens(ntokens);
718    r.set_vlen(vlen);
719    r.set_parser_state(ReqState::Start as u32);
720    r.set_parser_pos(next_pos);
721    r.set_parser_token(None);
722    r.set_parse_result(MsgParseResult::Ok);
723    MsgParseResult::Ok
724}
725
726#[allow(clippy::too_many_arguments)]
727fn finish_error(
728    r: &mut Msg,
729    state: ReqState,
730    pos: usize,
731    token: Option<usize>,
732    vlen: u32,
733    ty: MsgType,
734    is_read: bool,
735    quit: bool,
736    ntokens: u32,
737) -> MsgParseResult {
738    r.set_parser_state(state as u32);
739    r.set_parser_pos(pos);
740    r.set_parser_token(token);
741    r.set_vlen(vlen);
742    r.set_ntokens(ntokens);
743    if ty != MsgType::Unknown {
744        r.set_type(ty);
745    }
746    r.flags_mut().is_read = is_read;
747    r.flags_mut().quit = quit;
748    r.set_parse_result(MsgParseResult::Error);
749    MsgParseResult::Error
750}
751
752/// Parse a Memcached response from `input` and update `r` in place.
753///
754/// On success the response type is recorded and the parser cursor
755/// advances just past the trailing LF. The function never panics
756/// on any byte sequence.
757///
758/// # Examples
759///
760/// ```
761/// use dynomite::msg::{Msg, MsgParseResult, MsgType};
762/// use dynomite::proto::memcache::memcache_parse_rsp;
763///
764/// let mut r = Msg::new(0, MsgType::Unknown, false);
765/// let res = memcache_parse_rsp(&mut r, b"STORED\r\n");
766/// assert_eq!(res, MsgParseResult::Ok);
767/// assert_eq!(r.ty(), MsgType::RspMcStored);
768/// ```
769///
770/// The state machine intentionally lives in a single function so
771/// the byte-by-byte control flow stays in one place.
772#[allow(clippy::too_many_lines)]
773pub fn memcache_parse_rsp(r: &mut Msg, input: &[u8]) -> MsgParseResult {
774    if r.is_request() {
775        r.set_parse_result(MsgParseResult::Error);
776        return MsgParseResult::Error;
777    }
778    let mut state = RspState::from_u32(r.parser_state());
779    let mut p = r.parser_pos();
780    let mut token: Option<usize> = r.parser_token();
781    let mut vlen = r.vlen();
782    let mut ty = r.ty();
783    let mut end_marker = r.end_marker();
784
785    while p < input.len() {
786        let ch = input[p];
787        match state {
788            RspState::Start => {
789                if ch.is_ascii_digit() {
790                    state = RspState::RspNum;
791                } else {
792                    state = RspState::RspStr;
793                }
794                // Do not advance; re-enter under the new state.
795            }
796            RspState::RspNum => {
797                if token.is_none() {
798                    token = Some(p);
799                }
800                if ch.is_ascii_digit() {
801                    p += 1;
802                } else if ch == b' ' || ch == CR {
803                    token = None;
804                    ty = MsgType::RspMcNum;
805                    state = RspState::Crlf;
806                    // Do not advance.
807                } else {
808                    return finish_error_rsp(r, state, p, token, vlen, ty, end_marker);
809                }
810            }
811            RspState::RspStr => {
812                if token.is_none() {
813                    token = Some(p);
814                }
815                if ch == b' ' || ch == CR {
816                    let start = token.expect("token recorded");
817                    let key_bytes = &input[start..p];
818                    ty = classify_response(key_bytes);
819                    if ty == MsgType::RspMcEnd {
820                        end_marker = Some(start);
821                    }
822                    match ty {
823                        MsgType::Unknown => {
824                            return finish_error_rsp(r, state, p, token, vlen, ty, end_marker);
825                        }
826                        MsgType::RspMcStored
827                        | MsgType::RspMcNotStored
828                        | MsgType::RspMcExists
829                        | MsgType::RspMcNotFound
830                        | MsgType::RspMcDeleted
831                        | MsgType::RspMcTouched
832                        | MsgType::RspMcEnd
833                        | MsgType::RspMcError => {
834                            state = RspState::Crlf;
835                        }
836                        MsgType::RspMcValue => {
837                            state = RspState::SpacesBeforeKey;
838                        }
839                        MsgType::RspMcClientError | MsgType::RspMcServerError => {
840                            state = RspState::RuntoCrlf;
841                        }
842                        _ => {
843                            return finish_error_rsp(r, state, p, token, vlen, ty, end_marker);
844                        }
845                    }
846                    // Do not advance; re-enter on the same byte.
847                } else {
848                    p += 1;
849                }
850            }
851            RspState::SpacesBeforeKey => {
852                if ch == b' ' {
853                    p += 1;
854                } else {
855                    state = RspState::Key;
856                    // Do not advance.
857                }
858            }
859            RspState::Key => {
860                if ch == b' ' {
861                    state = RspState::SpacesBeforeFlags;
862                }
863                p += 1;
864            }
865            RspState::SpacesBeforeFlags => {
866                if ch == b' ' {
867                    p += 1;
868                } else if ch.is_ascii_digit() {
869                    state = RspState::Flags;
870                    // Do not advance.
871                } else {
872                    return finish_error_rsp(r, state, p, token, vlen, ty, end_marker);
873                }
874            }
875            RspState::Flags => {
876                if ch.is_ascii_digit() {
877                    p += 1;
878                } else if ch == b' ' {
879                    state = RspState::SpacesBeforeVlen;
880                    p += 1;
881                } else {
882                    return finish_error_rsp(r, state, p, token, vlen, ty, end_marker);
883                }
884            }
885            RspState::SpacesBeforeVlen => {
886                if ch == b' ' {
887                    p += 1;
888                } else if ch.is_ascii_digit() {
889                    state = RspState::Vlen;
890                    vlen = 0;
891                    // Do not advance.
892                } else {
893                    return finish_error_rsp(r, state, p, token, vlen, ty, end_marker);
894                }
895            }
896            RspState::Vlen => {
897                if ch.is_ascii_digit() {
898                    vlen = vlen.saturating_mul(10).saturating_add(u32::from(ch - b'0'));
899                    p += 1;
900                } else if ch == b' ' || ch == CR {
901                    state = RspState::RuntoCrlf;
902                    // Do not advance.
903                } else {
904                    return finish_error_rsp(r, state, p, token, vlen, ty, end_marker);
905                }
906            }
907            RspState::RuntoVal => match ch {
908                LF => {
909                    state = RspState::Val;
910                    token = None;
911                    p += 1;
912                }
913                _ => {
914                    return finish_error_rsp(r, state, p, token, vlen, ty, end_marker);
915                }
916            },
917            RspState::Val => {
918                let m = p.saturating_add(vlen as usize);
919                if m >= input.len() {
920                    let consumed = input.len() - p;
921                    vlen = vlen.saturating_sub(consumed as u32);
922                    p = input.len();
923                    break;
924                }
925                if input[m] != CR {
926                    return finish_error_rsp(r, state, p, token, vlen, ty, end_marker);
927                }
928                p = m + 1;
929                state = RspState::ValLf;
930            }
931            RspState::ValLf => match ch {
932                LF => {
933                    state = RspState::RspStr;
934                    p += 1;
935                }
936                _ => {
937                    return finish_error_rsp(r, state, p, token, vlen, ty, end_marker);
938                }
939            },
940            RspState::End => {
941                if token.is_none() {
942                    if ch != b'E' {
943                        return finish_error_rsp(r, state, p, token, vlen, ty, end_marker);
944                    }
945                    token = Some(p);
946                    p += 1;
947                } else if ch == CR {
948                    let start = token.expect("token recorded");
949                    if p - start == 3 && &input[start..p] == b"END" {
950                        end_marker = Some(start);
951                        state = RspState::AlmostDone;
952                        token = None;
953                        p += 1;
954                    } else {
955                        return finish_error_rsp(r, state, p, token, vlen, ty, end_marker);
956                    }
957                } else {
958                    p += 1;
959                }
960            }
961            RspState::RuntoCrlf => match ch {
962                CR => {
963                    if ty == MsgType::RspMcValue {
964                        state = RspState::RuntoVal;
965                    } else {
966                        state = RspState::AlmostDone;
967                    }
968                    p += 1;
969                }
970                _ => {
971                    p += 1;
972                }
973            },
974            RspState::Crlf => match ch {
975                b' ' => {
976                    p += 1;
977                }
978                CR => {
979                    state = RspState::AlmostDone;
980                    p += 1;
981                }
982                _ => {
983                    return finish_error_rsp(r, state, p, token, vlen, ty, end_marker);
984                }
985            },
986            RspState::AlmostDone => match ch {
987                LF => {
988                    r.set_type(ty);
989                    r.set_vlen(vlen);
990                    r.set_end_marker(end_marker);
991                    r.set_parser_state(RspState::Start as u32);
992                    r.set_parser_pos(p + 1);
993                    r.set_parser_token(None);
994                    r.set_parse_result(MsgParseResult::Ok);
995                    return MsgParseResult::Ok;
996                }
997                _ => {
998                    return finish_error_rsp(r, state, p, token, vlen, ty, end_marker);
999                }
1000            },
1001        }
1002    }
1003
1004    r.set_parser_state(state as u32);
1005    r.set_parser_pos(p);
1006    r.set_parser_token(token);
1007    r.set_vlen(vlen);
1008    r.set_end_marker(end_marker);
1009    if ty != MsgType::Unknown {
1010        r.set_type(ty);
1011    }
1012    r.set_parse_result(MsgParseResult::Again);
1013    MsgParseResult::Again
1014}
1015
1016fn finish_error_rsp(
1017    r: &mut Msg,
1018    state: RspState,
1019    pos: usize,
1020    token: Option<usize>,
1021    vlen: u32,
1022    ty: MsgType,
1023    end_marker: Option<usize>,
1024) -> MsgParseResult {
1025    r.set_parser_state(state as u32);
1026    r.set_parser_pos(pos);
1027    r.set_parser_token(token);
1028    r.set_vlen(vlen);
1029    r.set_end_marker(end_marker);
1030    if ty != MsgType::Unknown {
1031        r.set_type(ty);
1032    }
1033    r.set_parse_result(MsgParseResult::Error);
1034    MsgParseResult::Error
1035}
1036
1037#[cfg(test)]
1038mod tests {
1039    use super::*;
1040
1041    fn parse_req(input: &[u8]) -> Msg {
1042        let mut m = Msg::new(0, MsgType::Unknown, true);
1043        let _ = memcache_parse_req(&mut m, input);
1044        m
1045    }
1046
1047    fn parse_rsp(input: &[u8]) -> Msg {
1048        let mut m = Msg::new(0, MsgType::Unknown, false);
1049        let _ = memcache_parse_rsp(&mut m, input);
1050        m
1051    }
1052
1053    #[test]
1054    fn parse_get() {
1055        let m = parse_req(b"get key1\r\n");
1056        assert_eq!(m.parse_result(), MsgParseResult::Ok);
1057        assert_eq!(m.ty(), MsgType::ReqMcGet);
1058        assert_eq!(m.keys()[0].key(), b"key1");
1059        assert!(m.flags().is_read);
1060    }
1061
1062    #[test]
1063    fn parse_set() {
1064        let m = parse_req(b"set key1 0 0 3\r\nval\r\n");
1065        assert_eq!(m.parse_result(), MsgParseResult::Ok);
1066        assert_eq!(m.ty(), MsgType::ReqMcSet);
1067        assert_eq!(m.keys()[0].key(), b"key1");
1068        assert_eq!(m.vlen(), 3);
1069    }
1070
1071    #[test]
1072    fn parse_set_noreply() {
1073        let m = parse_req(b"set key1 0 0 3 noreply\r\nval\r\n");
1074        assert_eq!(m.parse_result(), MsgParseResult::Ok);
1075        assert!(!m.flags().expect_datastore_reply);
1076    }
1077
1078    #[test]
1079    fn parse_delete() {
1080        let m = parse_req(b"delete key1\r\n");
1081        assert_eq!(m.parse_result(), MsgParseResult::Ok);
1082        assert_eq!(m.ty(), MsgType::ReqMcDelete);
1083    }
1084
1085    #[test]
1086    fn parse_incr() {
1087        let m = parse_req(b"incr counter 1\r\n");
1088        assert_eq!(m.parse_result(), MsgParseResult::Ok);
1089        assert_eq!(m.ty(), MsgType::ReqMcIncr);
1090    }
1091
1092    #[test]
1093    fn parse_quit() {
1094        let m = parse_req(b"quit\r\n");
1095        assert_eq!(m.parse_result(), MsgParseResult::Ok);
1096        assert_eq!(m.ty(), MsgType::ReqMcQuit);
1097        assert!(m.flags().quit);
1098    }
1099
1100    #[test]
1101    fn parse_get_multikey() {
1102        let m = parse_req(b"get a b c\r\n");
1103        assert_eq!(m.parse_result(), MsgParseResult::Ok);
1104        let keys: Vec<&[u8]> = m.keys().iter().map(crate::msg::KeyPos::key).collect();
1105        assert_eq!(keys, vec![&b"a"[..], b"b", b"c"]);
1106    }
1107
1108    #[test]
1109    fn parse_cas() {
1110        let m = parse_req(b"cas key1 0 0 3 7\r\nval\r\n");
1111        assert_eq!(m.parse_result(), MsgParseResult::Ok);
1112        assert_eq!(m.ty(), MsgType::ReqMcCas);
1113    }
1114
1115    #[test]
1116    fn parse_too_long_key_errors() {
1117        let mut input = b"get ".to_vec();
1118        input.extend(std::iter::repeat_n(b'k', MEMCACHE_MAX_KEY_LENGTH + 1));
1119        input.extend_from_slice(b"\r\n");
1120        let m = parse_req(&input);
1121        assert_eq!(m.parse_result(), MsgParseResult::Error);
1122    }
1123
1124    #[test]
1125    fn parse_empty_key_errors() {
1126        let m = parse_req(b"get \r\n");
1127        assert_eq!(m.parse_result(), MsgParseResult::Error);
1128    }
1129
1130    #[test]
1131    fn parse_truncated_returns_again() {
1132        let m = parse_req(b"get key");
1133        assert_eq!(m.parse_result(), MsgParseResult::Again);
1134    }
1135
1136    #[test]
1137    fn parse_stored_response() {
1138        let m = parse_rsp(b"STORED\r\n");
1139        assert_eq!(m.ty(), MsgType::RspMcStored);
1140    }
1141
1142    #[test]
1143    fn parse_value_response() {
1144        let m = parse_rsp(b"VALUE key 0 3\r\nval\r\nEND\r\n");
1145        assert_eq!(m.parse_result(), MsgParseResult::Ok);
1146        assert_eq!(m.ty(), MsgType::RspMcEnd);
1147    }
1148
1149    #[test]
1150    fn parse_numeric_response() {
1151        let m = parse_rsp(b"42\r\n");
1152        assert_eq!(m.parse_result(), MsgParseResult::Ok);
1153        assert_eq!(m.ty(), MsgType::RspMcNum);
1154    }
1155
1156    #[test]
1157    fn parse_server_error_response() {
1158        let m = parse_rsp(b"SERVER_ERROR oops\r\n");
1159        assert_eq!(m.parse_result(), MsgParseResult::Ok);
1160        assert_eq!(m.ty(), MsgType::RspMcServerError);
1161    }
1162
1163    #[test]
1164    fn parse_response_unknown_keyword_errors() {
1165        let m = parse_rsp(b"BOGUS\r\n");
1166        assert_eq!(m.parse_result(), MsgParseResult::Error);
1167    }
1168}
dynomite/proto/memcache/parser.rs

dynomite/proto/memcache/
parser.rs
