dynomite/proto/redis/
verify.rs1use crate::msg::{Msg, MsgType};
10
11use super::fragment::FragmentDispatcher;
12
13pub const ADD_SET_STR: &[u8] = b"._add-set";
17pub const REM_SET_STR: &[u8] = b"._rem-set";
19
20pub fn redis_verify_request<D: FragmentDispatcher + ?Sized>(
44 r: &Msg,
45 dispatcher: &D,
46) -> Result<(), VerifyError> {
47 if r.ty() != MsgType::ReqRedisEval {
48 return Ok(());
49 }
50 if r.keys().len() <= 1 {
51 return Ok(());
52 }
53 let mut prev_idx: Option<u32> = None;
54 for k in r.keys() {
55 let key = k.key();
56 if key.starts_with(ADD_SET_STR) || key.starts_with(REM_SET_STR) {
57 continue;
58 }
59 let idx = dispatcher.shard_for(k.tag_bytes());
60 match prev_idx {
61 None => prev_idx = Some(idx),
62 Some(prev) if prev == idx => {}
63 Some(_) => return Err(VerifyError::ScriptSpansNodes),
64 }
65 }
66 Ok(())
67}
68
69#[derive(Copy, Clone, Debug, Eq, PartialEq, thiserror::Error)]
71#[non_exhaustive]
72pub enum VerifyError {
73 #[error("redis verify: script spans nodes")]
75 ScriptSpansNodes,
76}
77
78#[cfg(test)]
79mod tests {
80 use super::*;
81 use crate::msg::KeyPos;
82
83 struct OddEven;
84 impl FragmentDispatcher for OddEven {
85 fn shard_for(&self, key: &[u8]) -> u32 {
86 u32::from(*key.first().unwrap_or(&0)) % 2
87 }
88 fn shard_count(&self) -> u32 {
89 2
90 }
91 }
92
93 #[test]
94 fn non_eval_is_ok() {
95 let mut r = Msg::new(0, MsgType::ReqRedisGet, true);
96 r.push_key(KeyPos::without_tag(b"a".to_vec()));
97 r.push_key(KeyPos::without_tag(b"b".to_vec()));
98 assert!(redis_verify_request(&r, &OddEven).is_ok());
99 }
100
101 #[test]
102 fn eval_one_key_is_ok() {
103 let mut r = Msg::new(0, MsgType::ReqRedisEval, true);
104 r.push_key(KeyPos::without_tag(b"a".to_vec()));
105 assert!(redis_verify_request(&r, &OddEven).is_ok());
106 }
107
108 #[test]
109 fn eval_disjoint_shards_errors() {
110 let mut r = Msg::new(0, MsgType::ReqRedisEval, true);
111 r.push_key(KeyPos::without_tag(b"a".to_vec())); r.push_key(KeyPos::without_tag(b"b".to_vec())); assert_eq!(
114 redis_verify_request(&r, &OddEven),
115 Err(VerifyError::ScriptSpansNodes),
116 );
117 }
118
119 #[test]
120 fn eval_skips_metadata_keys() {
121 let mut r = Msg::new(0, MsgType::ReqRedisEval, true);
122 r.push_key(KeyPos::without_tag(b"a".to_vec()));
123 r.push_key(KeyPos::without_tag(b"._add-set".to_vec()));
124 assert!(redis_verify_request(&r, &OddEven).is_ok());
125 }
126}