From 812d07114cdd6541716fdc6edaf1d9ea2e71135e Mon Sep 17 00:00:00 2001
From: Gijs Burghoorn <g.burghoorn@gmail.com>
Date: Wed, 29 Mar 2023 20:25:20 +0200
Subject: [PATCH] Impr: Use signals to detect when X server is ready

---
 Cargo.lock          |  7 ++++
 Cargo.toml          |  3 ++
 src/post_login/x.rs | 94 ++++++++++++++++++++++++++++++---------------
 3 files changed, 72 insertions(+), 32 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 66845b2..68263f5 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -120,6 +120,7 @@ dependencies = [
  "libc 0.2.139",
  "log",
  "nix",
+ "once_cell",
  "pam",
  "pgs-files",
  "rand",
@@ -220,6 +221,12 @@ version = "0.1.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "b8f8bdf33df195859076e54ab11ee78a1b208382d3a26ec40d142ffc1ecc49ef"
 
+[[package]]
+name = "once_cell"
+version = "1.17.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b7e5500299e16ebb147ae15a00a942af264cf3688f47923b8fc2cd5858f23ad3"
+
 [[package]]
 name = "pam"
 version = "0.7.0"
diff --git a/Cargo.toml b/Cargo.toml
index a6cfed5..165be08 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -33,6 +33,9 @@ pam = "0.7.0"
 pgs-files = "0.0.7"
 users = "0.11.0"
 
+# Once Cell
+once_cell = "1.17.1"
+
 # Logging
 env_logger = { version = "0.9.0", default-features = false }
 log = "0.4.0"
diff --git a/src/post_login/x.rs b/src/post_login/x.rs
index e7e7a92..5436d39 100644
--- a/src/post_login/x.rs
+++ b/src/post_login/x.rs
@@ -1,4 +1,9 @@
+use libc::{signal, SIGUSR1, SIG_DFL, SIG_IGN};
 use rand::Rng;
+
+use once_cell::sync::Lazy;
+use std::sync::Mutex;
+
 use std::env;
 use std::error::Error;
 use std::fmt::Display;
@@ -13,7 +18,7 @@ use log::{error, info};
 use crate::auth::AuthUserInfo;
 use crate::env_container::EnvironmentContainer;
 
-const XSTART_TIMEOUT_SECS: u64 = 20;
+const XSTART_CHECK_MAX_TRIES: u64 = 300;
 const XSTART_CHECK_INTERVAL_MILLIS: u64 = 100;
 
 #[derive(Debug, Clone)]
@@ -51,6 +56,20 @@ fn mcookie() -> String {
     format!("{cookie:032x}")
 }
 
+static X_HAS_STARTED: Lazy<Mutex<bool>> = Lazy::new(|| Mutex::new(false));
+
+#[allow(dead_code)]
+fn handle_sigusr1(_: i32) {
+    *X_HAS_STARTED.lock().unwrap_or_else(|err| {
+        error!("Failed to grab the `X_HAS_STARTED` Mutex lock. Reason: {err}");
+        std::process::exit(1);
+    }) = true;
+
+    unsafe {
+        signal(SIGUSR1, handle_sigusr1 as usize);
+    }
+}
+
 pub fn setup_x(
     process_env: &mut EnvironmentContainer,
     user_info: &AuthUserInfo,
@@ -99,10 +118,20 @@ pub fn setup_x(
         vtnr_value
     };
 
+    // Here we explicitely ignore the first USR defined signal. Xorg looks at whether this signal
+    // is ignored or not. If it is ignored, it will send that signal to the parent when it ready to
+    // receive connections. This is also how xinit does it.
+    //
+    // After we spawn the Xorg process, we need to make sure to quickly re-enable this signal as we
+    // need to listen to the signal by Xorg.
+    unsafe {
+        libc::signal(SIGUSR1, SIG_IGN);
+    }
+
     info!("Run X server");
-    let child = Command::new(super::SYSTEM_SHELL)
+    let mut child = Command::new(super::SYSTEM_SHELL)
         .arg("-c")
-        .arg(format!("/usr/bin/X {display_value} vt{doubledigit_vtnr}",))
+        .arg(format!("/usr/bin/X {display_value} vt{doubledigit_vtnr}"))
         .stdout(Stdio::null()) // TODO: Maybe this should be logged or something?
         .stderr(Stdio::null()) // TODO: Maybe this should be logged or something?
         .spawn()
@@ -111,42 +140,43 @@ pub fn setup_x(
             XSetupError::XServerStart
         })?;
 
+    // See note above
+    unsafe {
+        libc::signal(SIGUSR1, SIG_DFL);
+        signal(SIGUSR1, handle_sigusr1 as usize);
+    }
+
     // Wait for XServer to boot-up
     let start_time = time::SystemTime::now();
-    loop {
-        // Timeout
-        if match start_time.elapsed() {
-            Ok(dur) => dur.as_secs() >= XSTART_TIMEOUT_SECS,
-            Err(_) => {
-                error!("Failed to resolve elapsed time");
-                std::process::exit(1);
-            }
-        } {
-            error!("Starting X timed out!");
-            return Err(XSetupError::XServerTimeout);
-        }
-
-        match Command::new(super::SYSTEM_SHELL)
-            .arg("-c")
-            .arg("timeout 1s /usr/bin/xset q")
-            .stdout(Stdio::null()) // TODO: Maybe this should be logged or something?
-            .stderr(Stdio::null()) // TODO: Maybe this should be logged or something?
-            .status()
-        {
-            Ok(status) => {
-                if status.success() {
-                    break;
-                }
-            }
-            Err(_) => {
-                error!("Failed to run `xset` to check X server status");
-                return Err(XSetupError::XServerStatusCheck);
-            }
+    for _ in 0..XSTART_CHECK_MAX_TRIES {
+        // This will be set by the `handle_sigusr1` signal handler.
+        if *X_HAS_STARTED.lock().unwrap() {
+            break;
         }
 
         thread::sleep(time::Duration::from_millis(XSTART_CHECK_INTERVAL_MILLIS));
     }
 
+    // If the value is still `false`, this means we have time-ed out and Xorg is not running.
+    if !*X_HAS_STARTED.lock().unwrap() {
+        child.kill().unwrap_or_else(|err| {
+            error!("Failed kill Xorg after it time-ed out. Reason: {err}");
+        });
+        return Err(XSetupError::XServerTimeout);
+    }
+
+    if let Ok(x_server_start_time) = start_time.elapsed() {
+        info!(
+            "It took X server {start_ms}ms to start",
+            start_ms = x_server_start_time.as_millis()
+        );
+    }
+
+    *X_HAS_STARTED.lock().unwrap_or_else(|err| {
+        error!("Failed to grab the `X_HAS_STARTED` Mutex lock. Reason: {err}");
+        std::process::exit(1);
+    }) = false;
+
     info!("X server is running");
 
     Ok(child)