1use std::io;
7
8use kevy_resp::Reply;
9
10#[non_exhaustive]
16#[derive(Debug, Clone, PartialEq, Eq)]
17pub enum PubsubEvent {
18 Subscribe {
20 channel: Vec<u8>,
22 count: i64,
24 },
25 Psubscribe {
27 pattern: Vec<u8>,
29 count: i64,
31 },
32 Unsubscribe {
34 channel: Option<Vec<u8>>,
36 count: i64,
38 },
39 Punsubscribe {
41 pattern: Option<Vec<u8>>,
43 count: i64,
45 },
46 Message {
48 channel: Vec<u8>,
50 payload: Vec<u8>,
52 },
53 Pmessage {
55 pattern: Vec<u8>,
57 channel: Vec<u8>,
59 payload: Vec<u8>,
61 },
62}
63
64pub fn classify_pubsub(reply: Reply) -> io::Result<PubsubEvent> {
68 let items = match reply {
69 Reply::Array(v) | Reply::Push(v) => v,
70 Reply::Error(e) => return Err(io::Error::other(String::from_utf8_lossy(&e).into_owned())),
71 other => {
72 return Err(invalid(format!(
73 "pubsub: expected array/push, got {}",
74 shape(&other)
75 )));
76 }
77 };
78
79 let mut it = items.into_iter();
80 let kind = take_bulk(
81 it.next().ok_or_else(|| invalid("pubsub: empty frame"))?,
82 "kind",
83 )?;
84
85 match kind.as_slice() {
86 b"subscribe" => {
87 let channel = take_bulk(
88 it.next().ok_or_else(|| invalid("subscribe: missing channel"))?,
89 "channel",
90 )?;
91 let count = take_int(
92 it.next().ok_or_else(|| invalid("subscribe: missing count"))?,
93 "count",
94 )?;
95 Ok(PubsubEvent::Subscribe { channel, count })
96 }
97 b"psubscribe" => {
98 let pattern = take_bulk(
99 it.next().ok_or_else(|| invalid("psubscribe: missing pattern"))?,
100 "pattern",
101 )?;
102 let count = take_int(
103 it.next().ok_or_else(|| invalid("psubscribe: missing count"))?,
104 "count",
105 )?;
106 Ok(PubsubEvent::Psubscribe { pattern, count })
107 }
108 b"unsubscribe" => {
109 let channel = take_bulk_or_nil(
110 it.next().ok_or_else(|| invalid("unsubscribe: missing channel"))?,
111 "channel",
112 )?;
113 let count = take_int(
114 it.next().ok_or_else(|| invalid("unsubscribe: missing count"))?,
115 "count",
116 )?;
117 Ok(PubsubEvent::Unsubscribe { channel, count })
118 }
119 b"punsubscribe" => {
120 let pattern = take_bulk_or_nil(
121 it.next()
122 .ok_or_else(|| invalid("punsubscribe: missing pattern"))?,
123 "pattern",
124 )?;
125 let count = take_int(
126 it.next()
127 .ok_or_else(|| invalid("punsubscribe: missing count"))?,
128 "count",
129 )?;
130 Ok(PubsubEvent::Punsubscribe { pattern, count })
131 }
132 b"message" => {
133 let channel = take_bulk(
134 it.next().ok_or_else(|| invalid("message: missing channel"))?,
135 "channel",
136 )?;
137 let payload = take_bulk(
138 it.next().ok_or_else(|| invalid("message: missing payload"))?,
139 "payload",
140 )?;
141 Ok(PubsubEvent::Message { channel, payload })
142 }
143 b"pmessage" => {
144 let pattern = take_bulk(
145 it.next().ok_or_else(|| invalid("pmessage: missing pattern"))?,
146 "pattern",
147 )?;
148 let channel = take_bulk(
149 it.next().ok_or_else(|| invalid("pmessage: missing channel"))?,
150 "channel",
151 )?;
152 let payload = take_bulk(
153 it.next().ok_or_else(|| invalid("pmessage: missing payload"))?,
154 "payload",
155 )?;
156 Ok(PubsubEvent::Pmessage {
157 pattern,
158 channel,
159 payload,
160 })
161 }
162 other => Err(invalid(format!(
163 "unknown pubsub kind: {}",
164 String::from_utf8_lossy(other)
165 ))),
166 }
167}
168
169fn take_bulk(r: Reply, field: &str) -> io::Result<Vec<u8>> {
170 match r {
171 Reply::Bulk(v) | Reply::Simple(v) => Ok(v),
172 other => Err(invalid(format!(
173 "pubsub field {field}: expected bulk, got {}",
174 shape(&other)
175 ))),
176 }
177}
178
179fn take_bulk_or_nil(r: Reply, field: &str) -> io::Result<Option<Vec<u8>>> {
180 match r {
181 Reply::Bulk(v) | Reply::Simple(v) => Ok(Some(v)),
182 Reply::Nil | Reply::Null => Ok(None),
183 other => Err(invalid(format!(
184 "pubsub field {field}: expected bulk/nil, got {}",
185 shape(&other)
186 ))),
187 }
188}
189
190fn take_int(r: Reply, field: &str) -> io::Result<i64> {
191 match r {
192 Reply::Int(n) => Ok(n),
193 other => Err(invalid(format!(
194 "pubsub field {field}: expected int, got {}",
195 shape(&other)
196 ))),
197 }
198}
199
200fn shape(r: &Reply) -> &'static str {
201 match r {
202 Reply::Simple(_) => "simple",
203 Reply::Error(_) => "error",
204 Reply::Int(_) => "int",
205 Reply::Bulk(_) => "bulk",
206 Reply::Nil | Reply::Null => "nil",
207 Reply::Array(_) => "array",
208 Reply::Map(_) => "map",
209 Reply::Set(_) => "set",
210 Reply::Double(_) => "double",
211 Reply::Boolean(_) => "boolean",
212 Reply::Verbatim { .. } => "verbatim",
213 Reply::BigNumber(_) => "bignumber",
214 Reply::Push(_) => "push",
215 Reply::BlobError(_) => "bloberror",
216 }
217}
218
219fn invalid(msg: impl Into<String>) -> io::Error {
220 io::Error::new(io::ErrorKind::InvalidData, msg.into())
221}
222
223#[cfg(test)]
224mod tests {
225 use super::*;
226
227 #[test]
228 fn classify_subscribe_ack() {
229 let r = Reply::Array(vec![
230 Reply::Bulk(b"subscribe".to_vec()),
231 Reply::Bulk(b"chan".to_vec()),
232 Reply::Int(1),
233 ]);
234 assert_eq!(
235 classify_pubsub(r).unwrap(),
236 PubsubEvent::Subscribe {
237 channel: b"chan".to_vec(),
238 count: 1,
239 }
240 );
241 }
242
243 #[test]
244 fn classify_message_event() {
245 let r = Reply::Array(vec![
246 Reply::Bulk(b"message".to_vec()),
247 Reply::Bulk(b"news".to_vec()),
248 Reply::Bulk(b"hello".to_vec()),
249 ]);
250 assert_eq!(
251 classify_pubsub(r).unwrap(),
252 PubsubEvent::Message {
253 channel: b"news".to_vec(),
254 payload: b"hello".to_vec(),
255 }
256 );
257 }
258
259 #[test]
260 fn classify_pmessage_event() {
261 let r = Reply::Array(vec![
262 Reply::Bulk(b"pmessage".to_vec()),
263 Reply::Bulk(b"news.*".to_vec()),
264 Reply::Bulk(b"news.tech".to_vec()),
265 Reply::Bulk(b"hi".to_vec()),
266 ]);
267 assert_eq!(
268 classify_pubsub(r).unwrap(),
269 PubsubEvent::Pmessage {
270 pattern: b"news.*".to_vec(),
271 channel: b"news.tech".to_vec(),
272 payload: b"hi".to_vec(),
273 }
274 );
275 }
276
277 #[test]
278 fn classify_unsubscribe_with_nil_channel() {
279 let r = Reply::Array(vec![
280 Reply::Bulk(b"unsubscribe".to_vec()),
281 Reply::Nil,
282 Reply::Int(0),
283 ]);
284 assert_eq!(
285 classify_pubsub(r).unwrap(),
286 PubsubEvent::Unsubscribe {
287 channel: None,
288 count: 0,
289 }
290 );
291 }
292
293 #[test]
294 fn classify_accepts_push_frame() {
295 let r = Reply::Push(vec![
297 Reply::Bulk(b"message".to_vec()),
298 Reply::Bulk(b"c".to_vec()),
299 Reply::Bulk(b"p".to_vec()),
300 ]);
301 assert_eq!(
302 classify_pubsub(r).unwrap(),
303 PubsubEvent::Message {
304 channel: b"c".to_vec(),
305 payload: b"p".to_vec(),
306 }
307 );
308 }
309
310 #[test]
311 fn classify_accepts_simple_string_fields() {
312 let r = Reply::Array(vec![
315 Reply::Simple(b"subscribe".to_vec()),
316 Reply::Simple(b"chan".to_vec()),
317 Reply::Int(2),
318 ]);
319 assert_eq!(
320 classify_pubsub(r).unwrap(),
321 PubsubEvent::Subscribe {
322 channel: b"chan".to_vec(),
323 count: 2,
324 }
325 );
326 }
327
328 #[test]
329 fn classify_rejects_unknown_kind() {
330 let r = Reply::Array(vec![
331 Reply::Bulk(b"bogus".to_vec()),
332 Reply::Bulk(b"x".to_vec()),
333 Reply::Int(0),
334 ]);
335 assert!(classify_pubsub(r).is_err());
336 }
337
338 #[test]
339 fn classify_rejects_wrong_arity() {
340 let r = Reply::Array(vec![
341 Reply::Bulk(b"subscribe".to_vec()),
342 Reply::Bulk(b"x".to_vec()),
343 ]);
344 assert!(classify_pubsub(r).is_err());
345 }
346}