From d03966ae422eaa35547b88e9c5ed97807eb78399 Mon Sep 17 00:00:00 2001
From: Lucy <lucy@absolucy.moe>
Date: Mon, 1 Aug 2022 13:18:24 -0400
Subject: [PATCH] Add integration to log out from COSMIC session

---
 .../cosmic-applet-power/src/cosmic_session.rs | 11 ++++++++++
 applets/cosmic-applet-power/src/main.rs       |  1 +
 applets/cosmic-applet-power/src/ui/session.rs | 20 ++++++++++++++++---
 3 files changed, 29 insertions(+), 3 deletions(-)
 create mode 100644 applets/cosmic-applet-power/src/cosmic_session.rs

diff --git a/applets/cosmic-applet-power/src/cosmic_session.rs b/applets/cosmic-applet-power/src/cosmic_session.rs
new file mode 100644
index 00000000..5d00dee5
--- /dev/null
+++ b/applets/cosmic-applet-power/src/cosmic_session.rs
@@ -0,0 +1,11 @@
+// SPDX-License-Identifier: GPL-3.0-or-later
+use zbus::dbus_proxy;
+
+#[dbus_proxy(
+    interface = "com.system76.CosmicSession",
+    default_service = "com.system76.CosmicSession",
+    default_path = "/com/system76/CosmicSession"
+)]
+trait CosmicSession {
+    fn exit(&self) -> zbus::Result<()>;
+}
diff --git a/applets/cosmic-applet-power/src/main.rs b/applets/cosmic-applet-power/src/main.rs
index 38b2bd5d..a4241d72 100644
--- a/applets/cosmic-applet-power/src/main.rs
+++ b/applets/cosmic-applet-power/src/main.rs
@@ -5,6 +5,7 @@ extern crate relm4_macros;
 
 pub mod session_manager;
 pub mod ui;
+pub mod cosmic_session;
 
 use gtk4::{gio::ApplicationFlags, prelude::*, Align, Button, Label, Orientation, Separator};
 use once_cell::sync::Lazy;
diff --git a/applets/cosmic-applet-power/src/ui/session.rs b/applets/cosmic-applet-power/src/ui/session.rs
index a2868b47..82bb69d1 100644
--- a/applets/cosmic-applet-power/src/ui/session.rs
+++ b/applets/cosmic-applet-power/src/ui/session.rs
@@ -1,6 +1,6 @@
 // SPDX-License-Identifier: GPL-3.0-or-later
 
-use crate::{session_manager::SessionManagerProxy, RT};
+use crate::{cosmic_session::CosmicSessionProxy, session_manager::SessionManagerProxy, RT};
 use gtk4::{prelude::*, Align, Button, Image, Label, Orientation};
 use logind_zbus::{
     manager::ManagerProxy,
@@ -35,9 +35,23 @@ async fn lock_screen() -> zbus::Result<()> {
 }
 
 async fn log_out() -> zbus::Result<()> {
+    let session_type = std::env::var("XDG_CURRENT_DESKTOP").ok();
     let connection = Connection::session().await?;
-    let manager_proxy = SessionManagerProxy::new(&connection).await?;
-    manager_proxy.logout(0).await
+    match session_type.as_ref().map(|s| s.trim()) {
+        Some("pop:COSMIC") => {
+            let cosmic_session = CosmicSessionProxy::new(&connection).await?;
+            cosmic_session.exit().await?;
+        }
+        Some("pop:GNOME") => {
+            let manager_proxy = SessionManagerProxy::new(&connection).await?;
+            manager_proxy.logout(0).await?;
+        }
+        Some(desktop) => {
+            eprintln!("unknown XDG_CURRENT_DESKTOP: {desktop}")
+        }
+        None => {}
+    }
+    Ok(())
 }
 
 pub fn build() -> gtk4::Box {