test: Add literal tests

flow-test/src/client_tester.rs

@@ -2,10 +2,12 @@ use std::net::SocketAddr;
 
 use bstr::ByteSlice;
 use imap_flow::{
-    client::{ClientFlow, ClientFlowError, ClientFlowEvent, ClientFlowOptions},
+    client::{
+        ClientFlow, ClientFlowCommandHandle, ClientFlowError, ClientFlowEvent, ClientFlowOptions,
+    },
     stream::AnyStream,
 };
-use imap_types::bounded_static::ToBoundedStatic;
+use imap_types::{bounded_static::ToBoundedStatic, command::Command};
 use tokio::net::TcpStream;
 use tracing::trace;
 
@@ -54,22 +56,62 @@ impl ClientTester {
         }
     }
 
-    pub async fn send_command(&mut self, bytes: &[u8]) {
-        let enqueued_command = self.codecs.decode_command_normalized(bytes);
+    pub fn enqueue_command(&mut self, bytes: &[u8]) -> EnqueuedCommand {
+        let command = self.codecs.decode_command_normalized(bytes).to_static();
+        let client = self.connection_state.greeted();
+        let handle = client.enqueue_command(command.to_static());
+        EnqueuedCommand { command, handle }
+    }
+
+    pub async fn progress_command(&mut self, enqueued_command: EnqueuedCommand) {
         let client = self.connection_state.greeted();
-        let enqueued_handle = client.enqueue_command(enqueued_command.to_static());
         let event = client.progress().await.unwrap();
         match event {
             ClientFlowEvent::CommandSent { handle, command } => {
-                assert_eq!(enqueued_handle, handle);
-                assert_eq!(enqueued_command, command);
+                assert_eq!(enqueued_command.handle, handle);
+                assert_eq!(enqueued_command.command, command);
+            }
+            event => {
+                panic!("Client emitted unexpected event: {event:?}");
+            }
+        }
+    }
+
+    pub async fn progress_rejected_command(
+        &mut self,
+        enqueued_command: EnqueuedCommand,
+        status_bytes: &[u8],
+    ) {
+        let expected_status = self.codecs.decode_status(status_bytes);
+        let client = self.connection_state.greeted();
+        let event = client.progress().await.unwrap();
+        match event {
+            ClientFlowEvent::CommandRejected {
+                handle,
+                command,
+                status,
+            } => {
+                assert_eq!(enqueued_command.handle, handle);
+                assert_eq!(enqueued_command.command, command);
+                assert_eq!(expected_status, status);
             }
             event => {
                 panic!("Client emitted unexpected event: {event:?}");
             }
         }
     }
 
+    pub async fn send_command(&mut self, bytes: &[u8]) {
+        let enqueued_command = self.enqueue_command(bytes);
+        self.progress_command(enqueued_command).await;
+    }
+
+    pub async fn send_rejected_command(&mut self, command_bytes: &[u8], status_bytes: &[u8]) {
+        let enqueued_command = self.enqueue_command(command_bytes);
+        self.progress_rejected_command(enqueued_command, status_bytes)
+            .await;
+    }
+
     pub async fn receive_data(&mut self, expected_bytes: &[u8]) {
         let expected_data = self.codecs.decode_data(expected_bytes);
         let client = self.connection_state.greeted();
@@ -153,3 +195,9 @@ impl ConnectionState {
         std::mem::replace(self, ConnectionState::Disconnected)
     }
 }
+
+/// A command that was enqueued and can later be used for assertions.
+pub struct EnqueuedCommand {
+    handle: ClientFlowCommandHandle,
+    command: Command<'static>,
+}

flow-test/src/runtime.rs

@@ -1,6 +1,6 @@
 use std::{future::Future, time::Duration};
 
-use tokio::{join, runtime, time::sleep};
+use tokio::{join, runtime, select, time::sleep};
 
 /// Options for creating an instance of `Runtime`.
 #[derive(Clone, Debug, PartialEq)]
@@ -60,4 +60,17 @@ impl Runtime {
     ) -> (T1, T2) {
         self.run(async { join!(future1, future2) })
     }
+
+    pub fn run2_and_select<T>(
+        &self,
+        future1: impl Future<Output = T>,
+        future2: impl Future<Output = T>,
+    ) -> T {
+        self.run(async {
+            select! {
+                output = future1 => output,
+                output = future2 => output,
+            }
+        })
+    }
 }

flow-test/src/server_tester.rs

@@ -112,6 +112,26 @@ impl ServerTester {
             }
         }
     }
+
+    pub async fn receive_error_because_literal_too_long(&mut self, expected_bytes: &[u8]) {
+        let server = self.connection_state.greeted();
+        let error = server.progress().await.unwrap_err();
+        match error {
+            ServerFlowError::LiteralTooLong { discarded_bytes } => {
+                assert_eq!(expected_bytes.as_bstr(), discarded_bytes.as_bstr());
+            }
+            error => {
+                panic!("Server has unexpected error: {error:?}");
+            }
+        }
+    }
+
+    /// Progresses internal responses without expecting any results.
+    pub async fn progress_internal_responses<T>(&mut self) -> T {
+        let server = self.connection_state.greeted();
+        let result = server.progress().await;
+        panic!("Server has unexpected result: {result:?}");
+    }
 }
 
 /// The current state of the connection between server and client.

flow-test/tests/both.rs

@@ -1,4 +1,5 @@
 use flow_test::test_setup::TestSetup;
+use imap_types::core::Text;
 
 #[test]
 fn noop() {
@@ -16,3 +17,58 @@ fn noop() {
     let status = b"A1 OK ...\r\n";
     rt.run2(server.send_status(status), client.receive_status(status));
 }
+
+#[test]
+fn login_with_literal() {
+    // The server will accept the literal ABCDE because it's smaller than the max size
+    let max_literal_size_tests = [5, 6, 10, 100];
+
+    for max_literal_size in max_literal_size_tests {
+        let mut setup = TestSetup::default();
+        setup.server_flow_options.literal_accept_text = Text::unvalidated("You shall pass");
+        setup.server_flow_options.max_literal_size = max_literal_size;
+
+        let (rt, mut server, mut client) = TestSetup::default().setup();
+
+        let greeting = b"* OK ...\r\n";
+        rt.run2(
+            server.send_greeting(greeting),
+            client.receive_greeting(greeting),
+        );
+
+        let login = b"A1 LOGIN {5}\r\nABCDE {5}\r\nFGHIJ\r\n";
+        rt.run2(client.send_command(login), server.receive_command(login));
+
+        let status = b"A1 NO ...\r\n";
+        rt.run2(server.send_status(status), client.receive_status(status));
+    }
+}
+
+#[test]
+fn login_with_rejected_literal() {
+    // The server will reject the literal ABCDE because it's larger than the max size
+    let max_literal_size_tests = [0, 1, 4];
+
+    for max_literal_size in max_literal_size_tests {
+        let mut setup = TestSetup::default();
+        setup.server_flow_options.literal_reject_text = Text::unvalidated("You shall not pass");
+        setup.server_flow_options.max_literal_size = max_literal_size;
+
+        let (rt, mut server, mut client) = setup.setup();
+
+        let greeting = b"* OK ...\r\n";
+        rt.run2(
+            server.send_greeting(greeting),
+            client.receive_greeting(greeting),
+        );
+
+        let login = b"A1 LOGIN {5}\r\nABCDE {5}\r\nFGHIJ\r\n";
+        let status = b"A1 BAD You shall not pass\r\n";
+        rt.run2_and_select(client.send_rejected_command(login, status), async {
+            server
+                .receive_error_because_literal_too_long(&login[..14])
+                .await;
+            server.progress_internal_responses().await
+        });
+    }
+}

flow-test/tests/client.rs

@@ -69,3 +69,80 @@ fn gibberish_instead_of_response() {
         client.receive_error_because_malformed_message(gibberish),
     );
 }
+
+#[test]
+fn login_with_literal() {
+    let (rt, mut server, mut client) = TestSetup::default().setup_client();
+
+    let greeting = b"* OK ...\r\n";
+    rt.run2(server.send(greeting), client.receive_greeting(greeting));
+
+    let login = b"A1 LOGIN {5}\r\nABCDE {5}\r\nFGHIJ\r\n";
+    let continuation_request = b"+ ...\r\n";
+    rt.run2(client.send_command(login), async {
+        server.receive(&login[..14]).await;
+        server.send(continuation_request).await;
+        server.receive(&login[14..25]).await;
+        server.send(continuation_request).await;
+        server.receive(&login[25..]).await;
+    });
+
+    let status = b"A1 NO ...\r\n";
+    rt.run2(server.send(status), client.receive_status(status));
+}
+
+#[test]
+fn login_with_rejected_literal() {
+    let (rt, mut server, mut client) = TestSetup::default().setup_client();
+
+    let greeting = b"* OK ...\r\n";
+    rt.run2(server.send(greeting), client.receive_greeting(greeting));
+
+    let login = b"A1 LOGIN {5}\r\nABCDE {5}\r\nFGHIJ\r\n";
+    let status = b"A1 BAD ...\r\n";
+    rt.run2(client.send_rejected_command(login, status), async {
+        server.receive(&login[..14]).await;
+        server.send(status).await;
+    });
+}
+
+#[test]
+fn login_with_literal_and_unexpected_status() {
+    // According to the specification, OK and NO will not affect the literal
+    let unexpected_status_tests = [b"A1 OK ...\r\n", b"A1 NO ...\r\n"];
+
+    for unexpected_status in unexpected_status_tests {
+        let (rt, mut server, mut client) = TestSetup::default().setup_client();
+
+        let greeting = b"* OK ...\r\n";
+        rt.run2(server.send(greeting), client.receive_greeting(greeting));
+
+        let login = b"A1 LOGIN {5}\r\nABCDE {5}\r\nFGHIJ\r\n";
+        let continuation_request = b"+ ...\r\n";
+        rt.run2(
+            async {
+                // Client starts sending the command
+                let command = client.enqueue_command(login);
+
+                // Client receives unexpected status
+                client.receive_status(unexpected_status).await;
+
+                // Client is able to continue sending the command
+                client.progress_command(command).await;
+            },
+            async {
+                // Server starts receiving the command
+                server.receive(&login[..14]).await;
+
+                // Server sends unexpected status
+                server.send(unexpected_status).await;
+
+                // Server continues receiving the command
+                server.send(continuation_request).await;
+                server.receive(&login[14..25]).await;
+                server.send(continuation_request).await;
+                server.receive(&login[25..]).await;
+            },
+        );
+    }
+}

flow-test/tests/server.rs

@@ -1,6 +1,7 @@
 use std::time::Duration;
 
 use flow_test::test_setup::TestSetup;
+use imap_types::core::Text;
 
 #[test]
 fn noop() {
@@ -55,3 +56,62 @@ fn gibberish_instead_of_command() {
         server.receive_error_because_malformed_message(gibberish),
     );
 }
+
+#[test]
+fn login_with_literal() {
+    // The server will accept the literal ABCDE because it's smaller than the max size
+    let max_literal_size_tests = [5, 6, 10, 100];
+
+    for max_literal_size in max_literal_size_tests {
+        let mut setup = TestSetup::default();
+        setup.server_flow_options.literal_accept_text = Text::unvalidated("You shall pass");
+        setup.server_flow_options.max_literal_size = max_literal_size;
+
+        let (rt, mut server, mut client) = setup.setup_server();
+
+        let greeting = b"* OK ...\r\n";
+        rt.run2(server.send_greeting(greeting), client.receive(greeting));
+
+        let login = b"A1 LOGIN {5}\r\nABCDE {5}\r\nFGHIJ\r\n";
+        let continuation_request = b"+ You shall pass\r\n";
+        rt.run2(
+            async {
+                client.send(&login[..14]).await;
+                client.receive(continuation_request).await;
+                client.send(&login[14..25]).await;
+                client.receive(continuation_request).await;
+                client.send(&login[25..]).await;
+            },
+            server.receive_command(login),
+        );
+
+        let status = b"A1 NO ...\r\n";
+        rt.run2(server.send_status(status), client.receive(status));
+    }
+}
+
+#[test]
+fn login_with_rejected_literal() {
+    // The server will reject the literal ABCDE because it's larger than the max size
+    let max_literal_size_tests = [0, 1, 4];
+
+    for max_literal_size in max_literal_size_tests {
+        let mut setup = TestSetup::default();
+        setup.server_flow_options.literal_reject_text = Text::unvalidated("You shall not pass");
+        setup.server_flow_options.max_literal_size = max_literal_size;
+
+        let (rt, mut server, mut client) = setup.setup_server();
+
+        let greeting = b"* OK ...\r\n";
+        rt.run2(server.send_greeting(greeting), client.receive(greeting));
+
+        let login = b"A1 LOGIN {5}\r\nABCDE {5}\r\nFGHIJ\r\n";
+        rt.run2(
+            client.send(&login[..14]),
+            server.receive_error_because_literal_too_long(&login[..14]),
+        );
+
+        let status = b"A1 BAD You shall not pass\r\n";
+        rt.run2_and_select(client.receive(status), server.progress_internal_responses());
+    }
+}