From 20ee03f4fd401f9b102c8d96c0d3f0f08e0379f1 Mon Sep 17 00:00:00 2001 From: Damian Poddebniak Date: Mon, 12 Feb 2024 16:15:38 +0100 Subject: [PATCH] test: Add literal tests --- flow-test/src/client_tester.rs | 62 +++++++++++++++++++++++---- flow-test/src/runtime.rs | 15 ++++++- flow-test/src/server_tester.rs | 20 +++++++++ flow-test/tests/both.rs | 56 +++++++++++++++++++++++++ flow-test/tests/client.rs | 77 ++++++++++++++++++++++++++++++++++ flow-test/tests/server.rs | 60 ++++++++++++++++++++++++++ 6 files changed, 282 insertions(+), 8 deletions(-) diff --git a/flow-test/src/client_tester.rs b/flow-test/src/client_tester.rs index b9f4a2d7..e3722bae 100644 --- a/flow-test/src/client_tester.rs +++ b/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,15 +56,44 @@ 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:?}"); @@ -70,6 +101,17 @@ impl ClientTester { } } + 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>, +} diff --git a/flow-test/src/runtime.rs b/flow-test/src/runtime.rs index 537b6191..721a6a76 100644 --- a/flow-test/src/runtime.rs +++ b/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( + &self, + future1: impl Future, + future2: impl Future, + ) -> T { + self.run(async { + select! { + output = future1 => output, + output = future2 => output, + } + }) + } } diff --git a/flow-test/src/server_tester.rs b/flow-test/src/server_tester.rs index a7aa69bf..44e82d6b 100644 --- a/flow-test/src/server_tester.rs +++ b/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(&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. diff --git a/flow-test/tests/both.rs b/flow-test/tests/both.rs index eb9c6036..d577aae9 100644 --- a/flow-test/tests/both.rs +++ b/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 + }); + } +} diff --git a/flow-test/tests/client.rs b/flow-test/tests/client.rs index bb5eddbe..9b0b5ec5 100644 --- a/flow-test/tests/client.rs +++ b/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; + }, + ); + } +} diff --git a/flow-test/tests/server.rs b/flow-test/tests/server.rs index 1b456dae..7af826d1 100644 --- a/flow-test/tests/server.rs +++ b/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()); + } +}