dynomite/proto/redis/
verify.rs1use crate::msg::{Msg, MsgType};
9
10use super::fragment::FragmentDispatcher;
11
12pub const ADD_SET_STR: &[u8] = b"._add-set";
16pub const REM_SET_STR: &[u8] = b"._rem-set";
18
19pub 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#[derive(Copy, Clone, Debug, Eq, PartialEq, thiserror::Error)]
70#[non_exhaustive]
71pub enum VerifyError {
72 #[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())); r.push_key(KeyPos::without_tag(b"b".to_vec())); 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 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 let mut r = Msg::new(0, MsgType::ReqRedisEval, true);
139 r.push_key(KeyPos::without_tag(b"a".to_vec())); r.push_key(KeyPos::without_tag(b"c".to_vec())); assert!(redis_verify_request(&r, &OddEven).is_ok());
142 }
143}