From dc92340774efd140ad17e306b9cbc2bd89740b9c Mon Sep 17 00:00:00 2001
From: Lucy <lucy@absolucy.moe>
Date: Mon, 21 Feb 2022 14:09:15 -0500
Subject: [PATCH] Locking the screen now works.

---
 Cargo.lock                                    |  1 +
 applets/cosmic-applet-power/Cargo.toml        |  1 +
 applets/cosmic-applet-power/src/ui/session.rs | 49 +++++++++++++++++++
 3 files changed, 51 insertions(+)

diff --git a/Cargo.lock b/Cargo.lock
index e3813ad1..83d37959 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -251,6 +251,7 @@ dependencies = [
  "futures-util",
  "gtk4",
  "logind-zbus",
+ "nix 0.23.1",
  "once_cell",
  "relm4-macros",
  "tokio",
diff --git a/applets/cosmic-applet-power/Cargo.toml b/applets/cosmic-applet-power/Cargo.toml
index e28922aa..3f8aa677 100644
--- a/applets/cosmic-applet-power/Cargo.toml
+++ b/applets/cosmic-applet-power/Cargo.toml
@@ -8,6 +8,7 @@ license = "LGPL-3.0-or-later"
 futures-util = "0.3.21"
 gtk4 = "0.4.6"
 logind-zbus = "3.0.1"
+nix = "0.23.1"
 once_cell = "1.9.0"
 relm4-macros = "0.4.1"
 tokio = { version = "1.15.0", features = ["full"] }
diff --git a/applets/cosmic-applet-power/src/ui/session.rs b/applets/cosmic-applet-power/src/ui/session.rs
index daa78997..0ce50260 100644
--- a/applets/cosmic-applet-power/src/ui/session.rs
+++ b/applets/cosmic-applet-power/src/ui/session.rs
@@ -1,6 +1,45 @@
 // SPDX-License-Identifier: LGPL-3.0-or-later
 
+use crate::RT;
 use gtk4::{prelude::*, Align, Button, Image, Label, Orientation};
+use logind_zbus::{
+    manager::ManagerProxy,
+    session::{SessionProxy, SessionType},
+    user::UserProxy,
+};
+use nix::unistd::getuid;
+use zbus::Connection;
+
+async fn lock_screen() -> zbus::Result<()> {
+    let connection = Connection::system().await?;
+    let manager_proxy = ManagerProxy::new(&connection).await?;
+    // Get the session this current process is running in
+    let our_uid = getuid().as_raw() as u32;
+    let user_path = manager_proxy.get_user(our_uid).await?;
+    let user = UserProxy::builder(&connection)
+        .path(user_path)?
+        .build()
+        .await?;
+    // Lock all non-TTY sessions of this user
+    let sessions = user.sessions().await?;
+    for (_, session_path) in sessions {
+        let session = SessionProxy::builder(&connection)
+            .path(session_path)?
+            .build()
+            .await?;
+        if session.type_().await? != SessionType::TTY {
+            session.lock().await?;
+        }
+    }
+    Ok(())
+}
+
+async fn log_out() -> zbus::Result<()> {
+    let connection = Connection::system().await?;
+    let manager_proxy = ManagerProxy::new(&connection).await?;
+    // TODO: figure out what function logs the current user out
+    manager_proxy.lock_sessions().await
+}
 
 pub fn build() -> gtk4::Box {
     view! {
@@ -45,5 +84,15 @@ pub fn build() -> gtk4::Box {
             }
         }
     }
+    lock_screen_button.connect_clicked(|_| {
+        RT.spawn(async move {
+            lock_screen().await.expect("failed to lock screen");
+        });
+    });
+    log_out_button.connect_clicked(|_| {
+        RT.spawn(async move {
+            log_out().await.expect("failed to log out");
+        });
+    });
     inner_box
 }