Add WaitForConnectFlow

Signed-off-by: Matthias Beyer <mail@beyermatthias.de>

1 files changed, 55 insertions, 0 deletions

diff --git a/mqtt-tester/src/flow.rs b/mqtt-tester/src/flow.rs
index 7706c19..83c3a06 100644
--- a/mqtt-tester/src/flow.rs
+++ b/mqtt-tester/src/flow.rs
@@ -17,3 +17,58 @@ pub trait Flow {
         command: crate::command::Command,
     ) -> Result<std::process::Output, miette::Error>;
 }
+
+pub struct WaitForConnectFlow;
+
+#[async_trait::async_trait]
+impl Flow for WaitForConnectFlow {
+    fn commands(&self) -> Vec<Box<dyn ClientExecutableCommand>> {
+        vec![Box::new(crate::executable::QuitCommand)]
+    }
+
+    async fn execute(
+        &self,
+        command: crate::command::Command,
+    ) -> Result<std::process::Output, miette::Error> {
+        command
+            .wait_for_write([crate::command::ClientCommand::WaitAndCheck(Box::new(
+                |bytes: &[u8]| -> bool {
+                    let connect_flags = if let Some(flags) = find_connect_flags(bytes) {
+                        flags
+                    } else {
+                        return false;
+                    };
+
+                    let username_flag_set = 0 != (connect_flags & 0b1000_0000); // Username flag
+                    let password_flag_set = 0 != (connect_flags & 0b0100_0000); // Username flag
+
+                    if username_flag_set {
+                        !password_flag_set
+                    } else {
+                        true
+                    }
+                },
+            ))])
+            .await
+    }
+}
+
+fn find_connect_flags(bytes: &[u8]) -> Option<u8> {
+    macro_rules! getbyte {
+        ($n:tt) => {
+            if let Some(b) = bytes.get($n) {
+                *b
+            } else {
+                return None;
+            }
+        };
+    }
+
+    if getbyte!(0) != 0b0001_0000 {
+        return None;
+    }
+
+    let str_len = getbyte!(4);
+    let connect_flag_position = 4usize + (str_len as usize) + 2;
+    Some(getbyte!(connect_flag_position))
+}