1use std::{fmt, sync::Arc, time::Duration};
2
3use bytes::Bytes;
4use rand::{Rng, RngCore};
5use tokio::sync::mpsc;
6
7#[cfg(test)]
8use crate::setup_tracing;
9
10use crate::{
11 TargetShard,
12 frame::{FrameOpcode, FrameParams, RequestFrame, RequestOpcode, ResponseFrame, ResponseOpcode},
13};
14use scylla_cql::Consistency;
15use scylla_cql::frame::protocol_features::ProtocolFeatures;
16use scylla_cql::frame::response::error::DbError;
17
18#[derive(Debug, Clone)]
22pub enum Condition {
23 True,
24
25 False,
26
27 Not(Box<Condition>),
28
29 And(Box<Condition>, Box<Condition>),
30
31 Or(Box<Condition>, Box<Condition>),
32
33 ConnectionSeqNo(usize),
35
36 RequestOpcode(RequestOpcode),
38
39 ResponseOpcode(ResponseOpcode),
41
42 BodyContainsCaseSensitive(Box<[u8]>),
44
45 BodyContainsCaseInsensitive(Box<[u8]>),
47
48 RandomWithProbability(f64),
50
51 TrueForLimitedTimes(usize),
53
54 ConnectionRegisteredAnyEvent,
56
57 RequestConsistency(Consistency, ProtocolFeatures),
63}
64
65pub(crate) struct EvaluationContext {
67 pub(crate) connection_seq_no: usize,
68 pub(crate) connection_has_events: bool,
69 pub(crate) opcode: FrameOpcode,
70 pub(crate) frame_body: Bytes,
71}
72
73impl Condition {
74 pub(crate) fn eval(&mut self, ctx: &EvaluationContext) -> bool {
75 match self {
76 Condition::True => true,
77
78 Condition::False => false,
79
80 Condition::Not(c) => !c.eval(ctx),
81
82 Condition::And(c1, c2) => c1.eval(ctx) && c2.eval(ctx),
83
84 Condition::Or(c1, c2) => c1.eval(ctx) || c2.eval(ctx),
85
86 Condition::ConnectionSeqNo(no) => *no == ctx.connection_seq_no,
87
88 Condition::RequestOpcode(op1) => match ctx.opcode {
89 FrameOpcode::Request(op2) => *op1 == op2,
90 FrameOpcode::Response(_) => panic!(
91 "Invalid type applied in rule condition: driver request opcode in cluster context"
92 ),
93 },
94
95 Condition::ResponseOpcode(op1) => match ctx.opcode {
96 FrameOpcode::Request(_) => panic!(
97 "Invalid type applied in rule condition: cluster response opcode in driver context"
98 ),
99 FrameOpcode::Response(op2) => *op1 == op2,
100 },
101
102 Condition::BodyContainsCaseSensitive(pattern) => ctx
103 .frame_body
104 .windows(pattern.len())
105 .any(|window| *window == **pattern),
106
107 Condition::BodyContainsCaseInsensitive(pattern) => std::str::from_utf8(pattern)
108 .map(|pattern_str| {
109 ctx.frame_body.windows(pattern.len()).any(|window| {
110 std::str::from_utf8(window)
111 .map(|window_str| str::eq_ignore_ascii_case(window_str, pattern_str))
112 .unwrap_or(false)
113 })
114 })
115 .unwrap_or(false),
116 Condition::RandomWithProbability(probability) => rand::rng().random_bool(*probability),
117
118 Condition::TrueForLimitedTimes(times) => {
119 let val = *times > 0;
120 if val {
121 *times -= 1;
122 }
123 val
124 }
125
126 Condition::ConnectionRegisteredAnyEvent => ctx.connection_has_events,
127
128 Condition::RequestConsistency(expected_cl, features) => match ctx.opcode {
129 FrameOpcode::Request(opcode) => {
130 let frame = RequestFrame::new(
131 FrameParams {
132 version: 0x04,
133 flags: 0,
134 stream: 0,
135 },
136 opcode,
137 ctx.frame_body.clone(),
138 );
139 frame
140 .deserialize(features)
141 .ok()
142 .and_then(|req| req.get_consistency())
143 == Some(*expected_cl)
144 }
145 FrameOpcode::Response(_) => false,
146 },
147 }
148 }
149
150 #[expect(clippy::should_implement_trait)]
152 pub fn not(c: Self) -> Self {
153 Condition::Not(Box::new(c))
154 }
155
156 pub fn and(self, c2: Self) -> Self {
158 Self::And(Box::new(self), Box::new(c2))
159 }
160
161 pub fn or(self, c2: Self) -> Self {
163 Self::Or(Box::new(self), Box::new(c2))
164 }
165
166 pub fn all(cs: impl IntoIterator<Item = Self>) -> Self {
168 let mut cs = cs.into_iter();
169 match cs.next() {
170 None => Self::True, Some(mut c) => {
172 for head in cs {
173 c = head.and(c);
174 }
175 c
176 }
177 }
178 }
179
180 pub fn any(cs: impl IntoIterator<Item = Self>) -> Self {
182 let mut cs = cs.into_iter();
183 match cs.next() {
184 None => Self::False, Some(mut c) => {
186 for head in cs {
187 c = head.or(c);
188 }
189 c
190 }
191 }
192 }
193}
194
195pub trait Reaction: Sized {
207 type Incoming;
208 type Returning;
209
210 fn noop() -> Self;
212
213 fn drop_frame() -> Self;
215
216 fn delay(time: Duration) -> Self;
218
219 fn forge_response(f: Arc<dyn Fn(Self::Incoming) -> Self::Returning + Send + Sync>) -> Self;
221
222 fn forge_response_with_delay(
224 time: Duration,
225 f: Arc<dyn Fn(Self::Incoming) -> Self::Returning + Send + Sync>,
226 ) -> Self;
227
228 fn transform_frame(f: Arc<dyn Fn(Self::Incoming) -> Self::Incoming + Send + Sync>) -> Self;
230
231 fn drop_connection() -> Self;
233
234 fn drop_connection_with_delay(time: Duration) -> Self;
236
237 fn with_feedback_when_performed(
240 self,
241 tx: mpsc::UnboundedSender<(Self::Incoming, Option<TargetShard>)>,
242 ) -> Self;
243}
244
245fn fmt_reaction(
246 f: &mut std::fmt::Formatter<'_>,
247 reaction_type: &str,
248 to_addressee: &dyn fmt::Debug,
249 to_sender: &dyn fmt::Debug,
250 drop_connection: &dyn fmt::Debug,
251 has_feedback_channel: bool,
252) -> std::fmt::Result {
253 f.debug_struct(reaction_type)
254 .field("to_addressee", to_addressee)
255 .field("to_sender", to_sender)
256 .field("drop_connection", drop_connection)
257 .field(
258 "feedback_channel",
259 if has_feedback_channel {
260 &"Some(<feedback_channel>)"
261 } else {
262 &"None"
263 },
264 )
265 .finish()
266}
267
268#[derive(Clone)]
269pub struct RequestReaction {
270 pub to_addressee: Option<Action<RequestFrame, RequestFrame>>,
271 pub to_sender: Option<Action<RequestFrame, ResponseFrame>>,
272 pub drop_connection: Option<Option<Duration>>,
273 pub feedback_channel: Option<mpsc::UnboundedSender<(RequestFrame, Option<TargetShard>)>>,
274}
275
276impl fmt::Debug for RequestReaction {
277 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
278 fmt_reaction(
279 f,
280 "RequestReaction",
281 &self.to_addressee,
282 &self.to_sender,
283 &self.drop_connection,
284 self.feedback_channel.is_some(),
285 )
286 }
287}
288
289#[derive(Clone)]
290pub struct ResponseReaction {
291 pub to_addressee: Option<Action<ResponseFrame, ResponseFrame>>,
292 pub to_sender: Option<Action<ResponseFrame, RequestFrame>>,
293 pub drop_connection: Option<Option<Duration>>,
294 pub feedback_channel: Option<mpsc::UnboundedSender<(ResponseFrame, Option<TargetShard>)>>,
295}
296
297impl fmt::Debug for ResponseReaction {
298 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
299 fmt_reaction(
300 f,
301 "ResponseReaction",
302 &self.to_addressee,
303 &self.to_sender,
304 &self.drop_connection,
305 self.feedback_channel.is_some(),
306 )
307 }
308}
309
310impl Reaction for RequestReaction {
311 type Incoming = RequestFrame;
312 type Returning = ResponseFrame;
313
314 fn noop() -> Self {
315 RequestReaction {
316 to_addressee: Some(Action {
317 delay: None,
318 msg_processor: None,
319 }),
320 to_sender: None,
321 drop_connection: None,
322 feedback_channel: None,
323 }
324 }
325
326 fn drop_frame() -> Self {
327 RequestReaction {
328 to_addressee: None,
329 to_sender: None,
330 drop_connection: None,
331 feedback_channel: None,
332 }
333 }
334
335 fn delay(time: Duration) -> Self {
336 RequestReaction {
337 to_addressee: Some(Action {
338 delay: Some(time),
339 msg_processor: None,
340 }),
341 to_sender: None,
342 drop_connection: None,
343 feedback_channel: None,
344 }
345 }
346
347 fn forge_response(f: Arc<dyn Fn(Self::Incoming) -> Self::Returning + Send + Sync>) -> Self {
348 RequestReaction {
349 to_addressee: None,
350 to_sender: Some(Action {
351 delay: None,
352 msg_processor: Some(f),
353 }),
354 drop_connection: None,
355 feedback_channel: None,
356 }
357 }
358
359 fn forge_response_with_delay(
360 time: Duration,
361 f: Arc<dyn Fn(Self::Incoming) -> Self::Returning + Send + Sync>,
362 ) -> Self {
363 RequestReaction {
364 to_addressee: None,
365 to_sender: Some(Action {
366 delay: Some(time),
367 msg_processor: Some(f),
368 }),
369 drop_connection: None,
370 feedback_channel: None,
371 }
372 }
373
374 fn transform_frame(f: Arc<dyn Fn(Self::Incoming) -> Self::Incoming + Send + Sync>) -> Self {
375 RequestReaction {
376 to_addressee: Some(Action {
377 delay: None,
378 msg_processor: Some(f),
379 }),
380 to_sender: None,
381 drop_connection: None,
382 feedback_channel: None,
383 }
384 }
385
386 fn drop_connection() -> Self {
387 RequestReaction {
388 to_addressee: None,
389 to_sender: None,
390 drop_connection: Some(None),
391 feedback_channel: None,
392 }
393 }
394
395 fn drop_connection_with_delay(time: Duration) -> Self {
396 RequestReaction {
397 to_addressee: None,
398 to_sender: None,
399 drop_connection: Some(Some(time)),
400 feedback_channel: None,
401 }
402 }
403
404 fn with_feedback_when_performed(
405 self,
406 tx: mpsc::UnboundedSender<(Self::Incoming, Option<TargetShard>)>,
407 ) -> Self {
408 Self {
409 feedback_channel: Some(tx),
410 ..self
411 }
412 }
413}
414
415impl RequestReaction {
416 pub fn forge_with_error_lazy(gen_error: Box<dyn Fn() -> DbError + Send + Sync>) -> Self {
417 Self::forge_with_error_lazy_delay(gen_error, None)
418 }
419 pub fn forge_with_error_lazy_delay(
422 gen_error: Box<dyn Fn() -> DbError + Send + Sync>,
423 delay: Option<Duration>,
424 ) -> Self {
425 RequestReaction {
426 to_addressee: None,
427 to_sender: Some(Action {
428 delay,
429 msg_processor: Some(Arc::new(move |request: RequestFrame| {
430 ResponseFrame::forged_error(request.params.for_response(), gen_error(), None)
431 .unwrap()
432 })),
433 }),
434 drop_connection: None,
435 feedback_channel: None,
436 }
437 }
438
439 pub fn forge_with_error(error: DbError) -> Self {
440 Self::forge_with_error_and_message(error, Some("Proxy-triggered error.".into()))
441 }
442
443 pub fn forge_with_error_and_message(error: DbError, msg: Option<String>) -> Self {
445 ResponseFrame::forged_error(
447 FrameParams {
448 version: 0,
449 flags: 0,
450 stream: 0,
451 },
452 error.clone(),
453 None,
454 )
455 .unwrap_or_else(|_| panic!("Invalid DbError provided: {error:#?}"));
456
457 RequestReaction {
458 to_addressee: None,
459 to_sender: Some(Action {
460 delay: None,
461 msg_processor: Some(Arc::new(move |request: RequestFrame| {
462 ResponseFrame::forged_error(
463 request.params.for_response(),
464 error.clone(),
465 msg.as_deref(),
466 )
467 .unwrap()
468 })),
469 }),
470 drop_connection: None,
471 feedback_channel: None,
472 }
473 }
474
475 pub fn forge() -> ResponseForger {
476 ResponseForger
477 }
478}
479
480pub mod example_db_errors {
481 use bytes::Bytes;
482 use scylla_cql::{
483 Consistency,
484 frame::response::error::{DbError, WriteType},
485 };
486
487 pub fn syntax_error() -> DbError {
488 DbError::SyntaxError
489 }
490 pub fn invalid() -> DbError {
491 DbError::Invalid
492 }
493 pub fn already_exists() -> DbError {
494 DbError::AlreadyExists {
495 keyspace: "proxy".into(),
496 table: "worker".into(),
497 }
498 }
499 pub fn function_failure() -> DbError {
500 DbError::FunctionFailure {
501 keyspace: "proxy".into(),
502 function: "fibonacci".into(),
503 arg_types: vec!["n".into()],
504 }
505 }
506 pub fn authentication_error() -> DbError {
507 DbError::AuthenticationError
508 }
509 pub fn unauthorized() -> DbError {
510 DbError::Unauthorized
511 }
512 pub fn config_error() -> DbError {
513 DbError::ConfigError
514 }
515 pub fn unavailable() -> DbError {
516 DbError::Unavailable {
517 consistency: Consistency::One,
518 required: 2,
519 alive: 1,
520 }
521 }
522 pub fn overloaded() -> DbError {
523 DbError::Overloaded
524 }
525 pub fn is_bootstrapping() -> DbError {
526 DbError::IsBootstrapping
527 }
528 pub fn truncate_error() -> DbError {
529 DbError::TruncateError
530 }
531 pub fn read_timeout() -> DbError {
532 DbError::ReadTimeout {
533 consistency: Consistency::One,
534 received: 2,
535 required: 3,
536 data_present: true,
537 }
538 }
539 pub fn write_timeout() -> DbError {
540 DbError::WriteTimeout {
541 consistency: Consistency::One,
542 received: 2,
543 required: 3,
544 write_type: WriteType::UnloggedBatch,
545 }
546 }
547 pub fn read_failure() -> DbError {
548 DbError::ReadFailure {
549 consistency: Consistency::One,
550 received: 2,
551 required: 3,
552 data_present: true,
553 numfailures: 1,
554 }
555 }
556 pub fn write_failure() -> DbError {
557 DbError::WriteFailure {
558 consistency: Consistency::One,
559 received: 2,
560 required: 3,
561 write_type: WriteType::UnloggedBatch,
562 numfailures: 1,
563 }
564 }
565 pub fn unprepared() -> DbError {
566 DbError::Unprepared {
567 statement_id: Bytes::from_static(b"21372137"),
568 }
569 }
570 pub fn server_error() -> DbError {
571 DbError::ServerError
572 }
573 pub fn protocol_error() -> DbError {
574 DbError::ProtocolError
575 }
576 pub fn other(num: i32) -> DbError {
577 DbError::Other(num)
578 }
579}
580
581pub struct ResponseForger;
582
583impl ResponseForger {
584 pub fn syntax_error(&self) -> RequestReaction {
585 RequestReaction::forge_with_error(example_db_errors::syntax_error())
586 }
587 pub fn invalid(&self) -> RequestReaction {
588 RequestReaction::forge_with_error(example_db_errors::invalid())
589 }
590 pub fn already_exists(&self) -> RequestReaction {
591 RequestReaction::forge_with_error(example_db_errors::already_exists())
592 }
593 pub fn function_failure(&self) -> RequestReaction {
594 RequestReaction::forge_with_error(example_db_errors::function_failure())
595 }
596 pub fn authentication_error(&self) -> RequestReaction {
597 RequestReaction::forge_with_error(example_db_errors::authentication_error())
598 }
599 pub fn unauthorized(&self) -> RequestReaction {
600 RequestReaction::forge_with_error(example_db_errors::unauthorized())
601 }
602 pub fn config_error(&self) -> RequestReaction {
603 RequestReaction::forge_with_error(example_db_errors::config_error())
604 }
605 pub fn unavailable(&self) -> RequestReaction {
606 RequestReaction::forge_with_error(example_db_errors::unavailable())
607 }
608 pub fn overloaded(&self) -> RequestReaction {
609 RequestReaction::forge_with_error(example_db_errors::overloaded())
610 }
611 pub fn is_bootstrapping(&self) -> RequestReaction {
612 RequestReaction::forge_with_error(example_db_errors::is_bootstrapping())
613 }
614 pub fn truncate_error(&self) -> RequestReaction {
615 RequestReaction::forge_with_error(example_db_errors::truncate_error())
616 }
617 pub fn read_timeout(&self) -> RequestReaction {
618 RequestReaction::forge_with_error(example_db_errors::read_timeout())
619 }
620 pub fn write_timeout(&self) -> RequestReaction {
621 RequestReaction::forge_with_error(example_db_errors::write_timeout())
622 }
623 pub fn read_failure(&self) -> RequestReaction {
624 RequestReaction::forge_with_error(example_db_errors::read_failure())
625 }
626 pub fn write_failure(&self) -> RequestReaction {
627 RequestReaction::forge_with_error(example_db_errors::write_failure())
628 }
629 pub fn unprepared(&self) -> RequestReaction {
630 RequestReaction::forge_with_error(example_db_errors::unprepared())
631 }
632 pub fn server_error(&self) -> RequestReaction {
633 RequestReaction::forge_with_error(example_db_errors::server_error())
634 }
635 pub fn protocol_error(&self) -> RequestReaction {
636 RequestReaction::forge_with_error(example_db_errors::protocol_error())
637 }
638 pub fn other(&self, num: i32) -> RequestReaction {
639 RequestReaction::forge_with_error(example_db_errors::other(num))
640 }
641 pub fn random_error(&self) -> RequestReaction {
642 self.random_error_with_delay(None)
643 }
644 pub fn random_error_with_delay(&self, delay: Option<Duration>) -> RequestReaction {
645 static ERRORS: &[fn() -> DbError] = &[
646 example_db_errors::invalid,
647 example_db_errors::already_exists,
648 example_db_errors::function_failure,
649 example_db_errors::authentication_error,
650 example_db_errors::unauthorized,
651 example_db_errors::config_error,
652 example_db_errors::unavailable,
653 example_db_errors::overloaded,
654 example_db_errors::is_bootstrapping,
655 example_db_errors::truncate_error,
656 example_db_errors::read_timeout,
657 example_db_errors::write_timeout,
658 example_db_errors::write_failure,
659 example_db_errors::unprepared,
660 example_db_errors::server_error,
661 example_db_errors::protocol_error,
662 || example_db_errors::other(2137),
663 ];
664 RequestReaction::forge_with_error_lazy_delay(
665 Box::new(|| ERRORS[rand::rng().next_u32() as usize % ERRORS.len()]()),
666 delay,
667 )
668 }
669}
670
671impl Reaction for ResponseReaction {
672 type Incoming = ResponseFrame;
673 type Returning = RequestFrame;
674
675 fn noop() -> Self {
676 ResponseReaction {
677 to_addressee: Some(Action {
678 delay: None,
679 msg_processor: None,
680 }),
681 to_sender: None,
682 drop_connection: None,
683 feedback_channel: None,
684 }
685 }
686
687 fn drop_frame() -> Self {
688 ResponseReaction {
689 to_addressee: None,
690 to_sender: None,
691 drop_connection: None,
692 feedback_channel: None,
693 }
694 }
695
696 fn delay(time: Duration) -> Self {
697 ResponseReaction {
698 to_addressee: Some(Action {
699 delay: Some(time),
700 msg_processor: None,
701 }),
702 to_sender: None,
703 drop_connection: None,
704 feedback_channel: None,
705 }
706 }
707
708 fn forge_response(f: Arc<dyn Fn(Self::Incoming) -> Self::Returning + Send + Sync>) -> Self {
709 ResponseReaction {
710 to_addressee: None,
711 to_sender: Some(Action {
712 delay: None,
713 msg_processor: Some(f),
714 }),
715 drop_connection: None,
716 feedback_channel: None,
717 }
718 }
719
720 fn forge_response_with_delay(
721 time: Duration,
722 f: Arc<dyn Fn(Self::Incoming) -> Self::Returning + Send + Sync>,
723 ) -> Self {
724 ResponseReaction {
725 to_addressee: None,
726 to_sender: Some(Action {
727 delay: Some(time),
728 msg_processor: Some(f),
729 }),
730 drop_connection: None,
731 feedback_channel: None,
732 }
733 }
734
735 fn transform_frame(f: Arc<dyn Fn(Self::Incoming) -> Self::Incoming + Send + Sync>) -> Self {
736 ResponseReaction {
737 to_addressee: Some(Action {
738 delay: None,
739 msg_processor: Some(f),
740 }),
741 to_sender: None,
742 drop_connection: None,
743 feedback_channel: None,
744 }
745 }
746
747 fn drop_connection() -> Self {
748 ResponseReaction {
749 to_addressee: None,
750 to_sender: None,
751 drop_connection: Some(None),
752 feedback_channel: None,
753 }
754 }
755
756 fn drop_connection_with_delay(time: Duration) -> Self {
757 ResponseReaction {
758 to_addressee: None,
759 to_sender: None,
760 drop_connection: Some(Some(time)),
761 feedback_channel: None,
762 }
763 }
764
765 fn with_feedback_when_performed(
766 self,
767 tx: mpsc::UnboundedSender<(Self::Incoming, Option<TargetShard>)>,
768 ) -> Self {
769 Self {
770 feedback_channel: Some(tx),
771 ..self
772 }
773 }
774}
775
776#[derive(Clone)]
779pub struct Action<TFrom, TTo> {
780 pub delay: Option<Duration>,
781 pub msg_processor: Option<Arc<dyn Fn(TFrom) -> TTo + Send + Sync>>,
782}
783
784impl<TFrom, TTo> std::fmt::Debug for Action<TFrom, TTo> {
787 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
788 f.debug_struct("Action")
789 .field("delay", &self.delay)
790 .field(
791 "msg_processor",
792 match self.msg_processor {
793 Some(_) => &"Some(<closure>)",
794 None => &"None",
795 },
796 )
797 .finish()
798 }
799}
800
801#[derive(Clone, Debug)]
804pub struct RequestRule(pub Condition, pub RequestReaction);
805
806#[derive(Clone, Debug)]
809pub struct ResponseRule(pub Condition, pub ResponseReaction);
810
811#[test]
812fn condition_case_insensitive_matching() {
813 setup_tracing();
814 let mut condition_matching =
815 Condition::BodyContainsCaseInsensitive(Box::new(*b"cassandra'sInefficiency"));
816 let mut condition_nonmatching =
817 Condition::BodyContainsCaseInsensitive(Box::new(*b"cassandrasInefficiency"));
818 let ctx = EvaluationContext {
819 connection_seq_no: 42,
820 opcode: FrameOpcode::Request(RequestOpcode::Options),
821 frame_body: Bytes::from_static(b"\0\0x{0x223}Cassandra'sINEFFICIENCY\x12\x31"),
822 connection_has_events: false,
823 };
824
825 assert!(condition_matching.eval(&ctx));
826 assert!(!condition_nonmatching.eval(&ctx));
827}