Skip to main content

dynomite/proto/redis/
verify.rs

1//! Redis request structural verifier.
2//!
3//! Most commands have nothing to verify here; the only live check
4//! is that an `EVAL` request keeps every key it touches on the
5//! same shard. The verifier consults a caller-supplied dispatcher
6//! so the cluster-state coupling stays out of the protocol layer.
7
8use crate::msg::{Msg, MsgType};
9
10use super::fragment::FragmentDispatcher;
11
12/// Reserved metadata key prefixes the verifier ignores. These match
13/// the `_add-set` / `_rem-set` constants used by the read-repair
14/// engine.
15pub const ADD_SET_STR: &[u8] = b"._add-set";
16/// See [`ADD_SET_STR`].
17pub const REM_SET_STR: &[u8] = b"._rem-set";
18
19/// Run the parser-level structural checks against `r`.
20///
21/// Returns `Ok(())` when no rule was violated. Returns
22/// [`VerifyError::ScriptSpansNodes`] when an `EVAL` request touches
23/// keys belonging to more than one shard.
24///
25/// # Examples
26///
27/// ```
28/// use dynomite::msg::{KeyPos, Msg, MsgType};
29/// use dynomite::proto::redis::{redis_verify_request, FragmentDispatcher};
30///
31/// struct AllOnOne;
32/// impl FragmentDispatcher for AllOnOne {
33///     fn shard_for(&self, _key: &[u8]) -> u32 { 0 }
34///     fn shard_count(&self) -> u32 { 1 }
35/// }
36///
37/// let mut r = Msg::new(0, MsgType::ReqRedisEval, true);
38/// r.push_key(KeyPos::without_tag(b"a".to_vec()));
39/// r.push_key(KeyPos::without_tag(b"b".to_vec()));
40/// assert!(redis_verify_request(&r, &AllOnOne).is_ok());
41/// ```
42pub fn redis_verify_request<D: FragmentDispatcher + ?Sized>(
43    r: &Msg,
44    dispatcher: &D,
45) -> Result<(), VerifyError> {
46    if r.ty() != MsgType::ReqRedisEval {
47        return Ok(());
48    }
49    if r.keys().len() <= 1 {
50        return Ok(());
51    }
52    let mut prev_idx: Option<u32> = None;
53    for k in r.keys() {
54        let key = k.key();
55        if key.starts_with(ADD_SET_STR) || key.starts_with(REM_SET_STR) {
56            continue;
57        }
58        let idx = dispatcher.shard_for(k.tag_bytes());
59        match prev_idx {
60            None => prev_idx = Some(idx),
61            Some(prev) if prev == idx => {}
62            Some(_) => return Err(VerifyError::ScriptSpansNodes),
63        }
64    }
65    Ok(())
66}
67
68/// Errors [`redis_verify_request`] can produce.
69#[derive(Copy, Clone, Debug, Eq, PartialEq, thiserror::Error)]
70#[non_exhaustive]
71pub enum VerifyError {
72    /// `EVAL` request touches keys on more than one shard.
73    #[error("redis verify: script spans nodes")]
74    ScriptSpansNodes,
75}
76
77#[cfg(test)]
78mod tests {
79    use super::*;
80    use crate::msg::KeyPos;
81
82    struct OddEven;
83    impl FragmentDispatcher for OddEven {
84        fn shard_for(&self, key: &[u8]) -> u32 {
85            u32::from(*key.first().unwrap_or(&0)) % 2
86        }
87        fn shard_count(&self) -> u32 {
88            2
89        }
90    }
91
92    #[test]
93    fn non_eval_is_ok() {
94        let mut r = Msg::new(0, MsgType::ReqRedisGet, true);
95        r.push_key(KeyPos::without_tag(b"a".to_vec()));
96        r.push_key(KeyPos::without_tag(b"b".to_vec()));
97        assert!(redis_verify_request(&r, &OddEven).is_ok());
98    }
99
100    #[test]
101    fn eval_one_key_is_ok() {
102        let mut r = Msg::new(0, MsgType::ReqRedisEval, true);
103        r.push_key(KeyPos::without_tag(b"a".to_vec()));
104        assert!(redis_verify_request(&r, &OddEven).is_ok());
105    }
106
107    #[test]
108    fn eval_disjoint_shards_errors() {
109        let mut r = Msg::new(0, MsgType::ReqRedisEval, true);
110        r.push_key(KeyPos::without_tag(b"a".to_vec())); // 0x61 -> shard 1
111        r.push_key(KeyPos::without_tag(b"b".to_vec())); // 0x62 -> shard 0
112        assert_eq!(
113            redis_verify_request(&r, &OddEven),
114            Err(VerifyError::ScriptSpansNodes),
115        );
116    }
117
118    #[test]
119    fn eval_skips_metadata_keys() {
120        let mut r = Msg::new(0, MsgType::ReqRedisEval, true);
121        r.push_key(KeyPos::without_tag(b"a".to_vec()));
122        r.push_key(KeyPos::without_tag(b"._add-set".to_vec()));
123        assert!(redis_verify_request(&r, &OddEven).is_ok());
124    }
125
126    #[test]
127    fn dispatcher_shard_count_is_reported() {
128        // Exercise the test dispatcher's shard_count accessor so the
129        // FragmentDispatcher contract is covered end to end.
130        assert_eq!(OddEven.shard_count(), 2);
131        assert_eq!(OddEven.shard_for(b"b"), 0);
132    }
133
134    #[test]
135    fn eval_two_keys_same_shard_is_ok() {
136        // Two EVAL keys that hash to the same shard pass the
137        // single-shard check (the prev == idx arm).
138        let mut r = Msg::new(0, MsgType::ReqRedisEval, true);
139        r.push_key(KeyPos::without_tag(b"a".to_vec())); // shard 1
140        r.push_key(KeyPos::without_tag(b"c".to_vec())); // 0x63 -> shard 1
141        assert!(redis_verify_request(&r, &OddEven).is_ok());
142    }
143}