From ec822e421f8e804d5e18a7dc516fde8c3347797a Mon Sep 17 00:00:00 2001
From: Jeremy Soller <jackpot51@gmail.com>
Date: Tue, 6 Feb 2024 09:25:48 -0700
Subject: [PATCH] Implement kiosk mode for cosmic-greeter

---
 src/main.rs    | 42 +++++++++++++++++++++++++++++++++++++++++-
 src/session.rs | 43 ++++++++++++++++++++++++-------------------
 2 files changed, 65 insertions(+), 20 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index 732998e4..435094bd 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -9,7 +9,7 @@ use smithay::{
 };
 
 use anyhow::{Context, Result};
-use std::{ffi::OsString, sync::Arc};
+use std::{env, ffi::OsString, process, sync::Arc};
 use tracing::{error, info, warn};
 
 use crate::wayland::handlers::compositor::client_compositor_state;
@@ -56,6 +56,18 @@ fn main() -> Result<()> {
     // potentially tell the session we are setup now
     session::setup_socket(event_loop.handle(), &state)?;
 
+    let mut args = env::args().skip(1);
+    let mut child_opt = if let Some(exec) = args.next() {
+        // Run command in kiosk mode
+        let mut command = process::Command::new(&exec);
+        command.args(args);
+        command.envs(session::get_env(&state)?);
+        info!("Running {:?}", exec);
+        Some(command.spawn()?)
+    } else {
+        None
+    };
+
     if let Err(err) = theme::watch_theme(event_loop.handle()) {
         warn!(?err, "Failed to watch theme");
     }
@@ -83,8 +95,36 @@ fn main() -> Result<()> {
 
         // send out events
         let _ = state.common.display_handle.flush_clients();
+
+        // check if kiosk child is running
+        if let Some(ref mut child) = child_opt {
+            match child.try_wait() {
+                // Kiosk child exited with status
+                Ok(Some(exit_status)) => {
+                    info!("Command exited with status {:?}", exit_status);
+                    match exit_status.code() {
+                        // Exiting with the same status as the kiosk child
+                        Some(code) => process::exit(code),
+                        // The kiosk child exited with signal, exiting with error
+                        None => process::exit(1),
+                    }
+                }
+                // Command still running
+                Ok(None) => {}
+                // Kiosk child disappeared, exiting with error
+                Err(err) => {
+                    warn!(?err, "Failed to wait for command");
+                    process::exit(1);
+                }
+            }
+        }
     })?;
 
+    // kill kiosk child if loop exited
+    if let Some(mut child) = child_opt {
+        let _ = child.kill();
+    }
+
     // drop eventloop & state before logger
     std::mem::drop(event_loop);
     std::mem::drop(state);
diff --git a/src/session.rs b/src/session.rs
index c6458c46..ddcff172 100644
--- a/src/session.rs
+++ b/src/session.rs
@@ -59,6 +59,29 @@ unsafe fn set_cloexec(fd: RawFd) -> rustix::io::Result<()> {
     rustix::io::fcntl_setfd(fd, flags | rustix::io::FdFlags::CLOEXEC)
 }
 
+pub fn get_env(state: &State) -> Result<HashMap<String, String>> {
+    let mut env = HashMap::new();
+    env.insert(
+        String::from("WAYLAND_DISPLAY"),
+        state
+            .common
+            .socket
+            .clone()
+            .into_string()
+            .map_err(|_| anyhow!("wayland socket is no valid utf-8 string?"))?,
+    );
+    if let Some(display) = state
+        .common
+        .shell
+        .xwayland_state
+        .as_ref()
+        .map(|s| s.display)
+    {
+        env.insert(String::from("DISPLAY"), format!(":{}", display));
+    }
+    Ok(env)
+}
+
 pub fn setup_socket(handle: LoopHandle<State>, state: &State) -> Result<()> {
     if let Ok(fd_num) = std::env::var("COSMIC_SESSION_SOCK") {
         if let Ok(fd) = fd_num.parse::<RawFd>() {
@@ -72,25 +95,7 @@ pub fn setup_socket(handle: LoopHandle<State>, state: &State) -> Result<()> {
                 }
             };
 
-            let mut env = HashMap::new();
-            env.insert(
-                String::from("WAYLAND_DISPLAY"),
-                state
-                    .common
-                    .socket
-                    .clone()
-                    .into_string()
-                    .map_err(|_| anyhow!("wayland socket is no valid utf-8 string?"))?,
-            );
-            if let Some(display) = state
-                .common
-                .shell
-                .xwayland_state
-                .as_ref()
-                .map(|s| s.display)
-            {
-                env.insert(String::from("DISPLAY"), format!(":{}", display));
-            }
+            let env = get_env(state)?;
             let message = serde_json::to_string(&Message::SetEnv { variables: env })
                 .with_context(|| "Failed to encode environment variables into json")?;
             let bytes = message.into_bytes();