From 3b86be0545b7b70a8c5c325b883f82a254b173ec Mon Sep 17 00:00:00 2001
From: Aaditya Dhruv <aadityadhruv@mailbox.org>
Date: Fri, 13 Oct 2023 00:58:59 -0500
Subject: [PATCH] Add new command for logging out of iamb session (#162)

---
 src/base.rs     |  1 +
 src/commands.rs | 18 ++++++++++++++++++
 src/main.rs     | 19 +++++++++++++++++--
 src/worker.rs   | 41 +++++++++++++++++++++++++++++++++++++++++
 4 files changed, 77 insertions(+), 2 deletions(-)

diff --git a/src/base.rs b/src/base.rs
index 07d825a..b70c42e 100644
--- a/src/base.rs
+++ b/src/base.rs
@@ -261,6 +261,7 @@ pub enum SendAction {
 pub enum HomeserverAction {
     /// Create a new room with an optional localpart.
     CreateRoom(Option<String>, CreateRoomType, CreateRoomFlags),
+    Logout(String, bool),
 }
 
 /// An action that the main program loop should.
diff --git a/src/commands.rs b/src/commands.rs
index 58a954f..6024239 100644
--- a/src/commands.rs
+++ b/src/commands.rs
@@ -471,6 +471,19 @@ fn iamb_open(desc: CommandDescription, ctx: &mut ProgContext) -> ProgResult {
     return Ok(step);
 }
 
+fn iamb_logout(desc: CommandDescription, ctx: &mut ProgContext) -> ProgResult {
+    let args = desc.arg.strings()?;
+
+    if args.len() != 1 {
+        return Result::Err(CommandError::InvalidArgument);
+    }
+
+    let iact = IambAction::from(HomeserverAction::Logout(args[0].clone(), desc.bang));
+    let step = CommandStep::Continue(iact.into(), ctx.context.take());
+
+    return Ok(step);
+}
+
 fn add_iamb_commands(cmds: &mut ProgramCommands) {
     cmds.add_command(ProgramCommand {
         name: "cancel".into(),
@@ -557,6 +570,11 @@ fn add_iamb_commands(cmds: &mut ProgramCommands) {
         aliases: vec![],
         f: iamb_editor,
     });
+    cmds.add_command(ProgramCommand {
+        name: "logout".into(),
+        aliases: vec![],
+        f: iamb_logout,
+    });
 }
 
 /// Initialize the default command state.
diff --git a/src/main.rs b/src/main.rs
index dfb3997..f5b8f5e 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -115,12 +115,12 @@ use modalkit::{
             WindowAction,
             WindowContainer,
         },
-        base::{MoveDir1D, OpenTarget, RepeatType},
+        base::{CloseFlags, MoveDir1D, OpenTarget, RepeatType, TabTarget},
         context::Resolve,
         key::KeyManager,
         store::Store,
     },
-    input::{bindings::BindingMachine, dialog::Pager, key::TerminalKey},
+    input::{bindings::BindingMachine, dialog::Pager, dialog::PromptYesNo, key::TerminalKey},
     widgets::{
         cmdbar::CommandBarState,
         screen::{FocusList, Screen, ScreenState, TabLayoutDescription},
@@ -571,6 +571,21 @@ impl Application {
 
                 Ok(vec![(action.into(), ctx)])
             },
+            HomeserverAction::Logout(user, true) => {
+                self.worker.logout(user)?;
+                let flags = CloseFlags::QUIT | CloseFlags::FORCE;
+                let act = TabAction::Close(TabTarget::All, flags);
+
+                Ok(vec![(act.into(), ctx)])
+            },
+            HomeserverAction::Logout(user, false) => {
+                let msg = "Would you like to logout?";
+                let act = IambAction::from(HomeserverAction::Logout(user, true));
+                let prompt = PromptYesNo::new(msg, vec![Action::from(act)]);
+                let prompt = Box::new(prompt);
+
+                Err(UIError::NeedConfirm(prompt))
+            },
         }
     }
 
diff --git a/src/worker.rs b/src/worker.rs
index 255d41c..0b602ee 100644
--- a/src/worker.rs
+++ b/src/worker.rs
@@ -455,6 +455,7 @@ pub type FetchedRoom = (MatrixRoom, DisplayName, Option<Tags>);
 pub enum WorkerTask {
     Init(AsyncProgramStore, ClientReply<()>),
     Login(LoginStyle, ClientReply<IambResult<EditInfo>>),
+    Logout(String, ClientReply<IambResult<EditInfo>>),
     GetInviter(Invited, ClientReply<IambResult<Option<RoomMember>>>),
     GetRoom(OwnedRoomId, ClientReply<IambResult<FetchedRoom>>),
     JoinRoom(String, ClientReply<IambResult<OwnedRoomId>>),
@@ -480,6 +481,9 @@ impl Debug for WorkerTask {
                     .field(&format_args!("_"))
                     .finish()
             },
+            WorkerTask::Logout(user_id, _) => {
+                f.debug_tuple("WorkerTask::Logout").field(user_id).finish()
+            },
             WorkerTask::GetInviter(invite, _) => {
                 f.debug_tuple("WorkerTask::GetInviter").field(invite).finish()
             },
@@ -550,6 +554,14 @@ impl Requester {
         return response.recv();
     }
 
+    pub fn logout(&self, user_id: String) -> IambResult<EditInfo> {
+        let (reply, response) = oneshot();
+
+        self.tx.send(WorkerTask::Logout(user_id, reply)).unwrap();
+
+        return response.recv();
+    }
+
     pub fn get_inviter(&self, invite: Invited) -> IambResult<Option<RoomMember>> {
         let (reply, response) = oneshot();
 
@@ -704,6 +716,10 @@ impl ClientWorker {
                 assert!(self.initialized);
                 reply.send(self.login_and_sync(style).await);
             },
+            WorkerTask::Logout(user_id, reply) => {
+                assert!(self.initialized);
+                reply.send(self.logout(user_id).await);
+            },
             WorkerTask::Members(room_id, reply) => {
                 assert!(self.initialized);
                 reply.send(self.members(room_id).await);
@@ -1073,6 +1089,31 @@ impl ClientWorker {
         Ok(Some(InfoMessage::from("Successfully logged in!")))
     }
 
+    async fn logout(&mut self, user_id: String) -> IambResult<EditInfo> {
+        // Verify that the user is logging out of the correct profile.
+        let curr = self.settings.profile.user_id.as_ref();
+
+        if user_id != curr {
+            let msg = format!("Incorrect user ID (currently logged in as {curr})");
+            let err = UIError::Failure(msg);
+
+            return Err(err);
+        }
+
+        // Send the logout request.
+        if let Err(e) = self.client.logout().await {
+            let msg = format!("Failed to logout: {e}");
+            let err = UIError::Failure(msg);
+
+            return Err(err);
+        }
+
+        // Remove the session.json file.
+        std::fs::remove_file(&self.settings.session_json)?;
+
+        Ok(Some(InfoMessage::from("Sucessfully logged out")))
+    }
+
     async fn direct_message(&mut self, user: OwnedUserId) -> IambResult<OwnedRoomId> {
         for room in self.client.rooms() {
             if !room.is_direct() {