feat: Implement AUTHENTICATE flow

duesee · Nov 14, 2023 · 9542a4c · 9542a4c
1 parent cf1143f
commit 9542a4c
Show file tree

Hide file tree

Showing 4 changed files with 124 additions and 11 deletions.
diff --git a/examples/client.rs b/examples/client.rs
@@ -1,8 +1,11 @@
 use imap_codec::imap_types::{
     command::{Command, CommandBody},
     core::Tag,
+    response::Status,
+    secret::Secret,
 };
 use imap_flow::{
+    auth::Authenticate,
     client::{ClientFlow, ClientFlowEvent, ClientFlowOptions},
     stream::AnyStream,
 };
@@ -18,9 +21,10 @@ async fn main() {
             .unwrap();
     println!("received greeting: {greeting:?}");
 
-    let handle = client.enqueue_command(Command {
+    // TODO: Forbid Command::Authenticate.
+    let handle1 = client.enqueue_command(Command {
         tag: Tag::try_from("A1").unwrap(),
-        body: CommandBody::Noop,
+        body: CommandBody::Capability,
     });
 
     loop {
@@ -30,22 +34,68 @@ async fn main() {
                 tag,
             } => {
                 println!("command sent: {got_handle:?}, {tag:?}");
-                assert_eq!(handle, got_handle);
+                assert_eq!(handle1, got_handle);
             }
             ClientFlowEvent::CommandRejected {
                 handle: got_handle,
                 tag,
                 status,
             } => {
                 println!("command rejected: {got_handle:?}, {tag:?}, {status:?}");
-                assert_eq!(handle, got_handle);
+                assert_eq!(handle1, got_handle);
             }
             ClientFlowEvent::DataReceived { data } => {
                 println!("data received: {data:?}");
             }
-            ClientFlowEvent::StatusReceived { status } => {
-                println!("status received: {status:?}");
+            ClientFlowEvent::StatusReceived { status } => match status {
+                Status::Tagged(_) => {
+                    println!("command completion result response received: {status:?}");
+                    break;
+                }
+                _ => {
+                    println!("status received: {status:?}");
+                }
+            },
+        }
+    }
+
+    let handle2 = client.enqueue_command_authenticate(
+        Tag::try_from("A2").unwrap(),
+        Authenticate::Plain {
+            username: b"alice".to_vec(),
+            password: Secret::new(b"password".to_vec()),
+        },
+    );
+
+    loop {
+        match client.progress().await.unwrap() {
+            ClientFlowEvent::CommandSent {
+                handle: got_handle,
+                tag,
+            } => {
+                println!("command sent: {got_handle:?}, {tag:?}");
+                assert_eq!(handle2, got_handle);
+            }
+            ClientFlowEvent::CommandRejected {
+                handle: got_handle,
+                tag,
+                status,
+            } => {
+                println!("command rejected: {got_handle:?}, {tag:?}, {status:?}");
+                assert_eq!(handle2, got_handle);
+            }
+            ClientFlowEvent::DataReceived { data } => {
+                println!("data received: {data:?}");
             }
+            ClientFlowEvent::StatusReceived { status } => match status {
+                Status::Tagged(_) => {
+                    println!("command completion result response received: {status:?}");
+                    break;
+                }
+                _ => {
+                    println!("status received: {status:?}");
+                }
+            },
         }
     }
 }
diff --git a/src/client.rs b/src/client.rs
@@ -3,15 +3,18 @@ use bytes::BytesMut;
 use imap_codec::{
     decode::{GreetingDecodeError, ResponseDecodeError},
     imap_types::{
-        command::Command,
+        auth::{AuthMechanism, AuthenticateData},
+        command::{Command, CommandBody},
         core::Tag,
         response::{Data, Greeting, Response, Status, StatusBody, StatusKind, Tagged},
+        secret::Secret,
     },
     CommandCodec, GreetingCodec, ResponseCodec,
 };
 use thiserror::Error;
 
 use crate::{
+    auth::Authenticate,
     receive::{ReceiveEvent, ReceiveState},
     send::SendCommandState,
     stream::AnyStream,
@@ -96,6 +99,40 @@ impl ClientFlow {
         handle
     }
 
+    pub fn enqueue_command_authenticate(
+        &mut self,
+        tag: Tag<'static>,
+        auth: Authenticate,
+    ) -> ClientFlowCommandHandle {
+        // TODO(#53)
+        let handle = self.next_command_handle;
+        self.next_command_handle = ClientFlowCommandHandle(handle.0 + 1);
+
+        let (cmd, data) = match auth {
+            Authenticate::Plain { username, password } => {
+                // SASL PLAIN has the following form:
+                //     base64(b"<authorization identity>\x00<authentication identity>\x00<password>")
+                // We don't use `<authorization identity>` (yet).
+                let mut data: Vec<u8> = vec![0x00];
+                data.extend_from_slice(&username);
+                data.push(0x00);
+                data.extend_from_slice(password.declassify());
+
+                (
+                    Command {
+                        tag: tag.clone(),
+                        body: CommandBody::authenticate(AuthMechanism::Plain),
+                    },
+                    AuthenticateData(Secret::new(data)),
+                )
+            }
+        };
+
+        self.send_command_state
+            .enqueue_authenticate((tag, handle), cmd, data);
+        handle
+    }
+
     pub async fn progress(&mut self) -> Result<ClientFlowEvent, ClientFlowError> {
         loop {
             if let Some(event) = self.progress_command().await? {

diff --git a/src/lib.rs b/src/lib.rs
@@ -5,3 +5,15 @@ mod receive;
 mod send;
 pub mod server;
 pub mod stream;
+
+pub mod auth {
+    use imap_codec::imap_types::secret::Secret;
+
+    #[derive(Debug)]
+    pub enum Authenticate {
+        Plain {
+            username: Vec<u8>,
+            password: Secret<Vec<u8>>,
+        },
+    }
+}
diff --git a/src/send.rs b/src/send.rs
@@ -3,8 +3,8 @@ use std::collections::VecDeque;
 use bytes::BytesMut;
 use imap_codec::{
     encode::{Encoder, Fragment},
-    imap_types::command::Command,
-    CommandCodec,
+    imap_types::{auth::AuthenticateData, command::Command},
+    AuthenticateDataCodec, CommandCodec,
 };
 use tokio::io::AsyncWriteExt;
 
@@ -38,6 +38,15 @@ impl<K> SendCommandState<K> {
         self.send_queue.push_back(entry);
     }
 
+    pub fn enqueue_authenticate(&mut self, key: K, command: Command<'_>, data: AuthenticateData) {
+        let mut fragments: VecDeque<_> = self.codec.encode(&command).collect();
+        let fragments_auth = AuthenticateDataCodec::new().encode(&data);
+        fragments.extend(fragments_auth);
+
+        let entry = SendCommandQueueEntry { key, fragments };
+        self.send_queue.push_back(entry);
+    }
+
     pub fn command_in_progress(&self) -> Option<&K> {
         self.send_progress.as_ref().map(|x| &x.key)
     }
@@ -122,8 +131,13 @@ impl<K> SendCommandState<K> {
                         });
                         break true;
                     }
-                    Fragment::AuthData { .. } => {
-                        unimplemented!()
+                    Fragment::AuthData { data } => {
+                        // Delay authentication data until `Continue` from server
+                        progress.next_literal = Some(SendCommandLiteralProgress {
+                            data,
+                            received_continue: false,
+                        });
+                        break true;
                     }
                 }
             } else {