Fix restart and exit commands not returning replies

locker/src/main.rs

@@ -167,7 +167,7 @@ fn main() -> anyhow::Result<()> {
         if let Err(err) = handle_xcb_events(&conn, &mut state, &command_atoms) {
             if conn.has_error().is_err() {
                 error!("Lost connection to X server; attempting to restart");
-                restart_daemon(&mut state)?;
+                restart_daemon(&mut state, None)?;
             }
             warn!("Error handling event: {}", err);
         }
@@ -211,7 +211,7 @@ fn main() -> anyhow::Result<()> {
                     if let Err(err) = result {
                         if conn.has_error().is_err() {
                             error!("Lost connection to X server; atempting to restart");
-                            restart_daemon(&mut state)?;
+                            restart_daemon(&mut state, None)?;
                         }
                         warn!("Error handling event: {}", err);
                     }
@@ -489,10 +489,10 @@ where
 fn handle_signals(state: &mut State, signal_fd: &mut SignalFd) -> anyhow::Result<()> {
     match signal_fd.read_signal()? {
         None => (),
-        Some(info) if info.ssi_signo == Signal::SIGHUP as u32 => restart_daemon(state)?,
+        Some(info) if info.ssi_signo == Signal::SIGHUP as u32 => restart_daemon(state, None)?,
-        Some(info) if info.ssi_signo == Signal::SIGINT as u32 => exit_daemon(state)?,
+        Some(info) if info.ssi_signo == Signal::SIGINT as u32 => exit_daemon(state, None)?,
-        Some(info) if info.ssi_signo == Signal::SIGQUIT as u32 => exit_daemon(state)?,
+        Some(info) if info.ssi_signo == Signal::SIGQUIT as u32 => exit_daemon(state, None)?,
-        Some(info) if info.ssi_signo == Signal::SIGTERM as u32 => exit_daemon(state)?,
+        Some(info) if info.ssi_signo == Signal::SIGTERM as u32 => exit_daemon(state, None)?,
         Some(info) => trace!("Unexpected signal {}", info.ssi_signo),
     }
     Ok(())
@@ -567,8 +567,8 @@ fn handle_xcb_events<'a>(conn: &'a xcb::Connection, state: &mut State<'a>, comma
                                     },
                                 }
                             },
-                            r if r == command_atoms.restart => Some(restart_daemon(state)),
+                            r if r == command_atoms.restart => Some(restart_daemon(state, Some((conn, &ev)))),
-                            e if e == command_atoms.exit => Some(exit_daemon(state)),
+                            e if e == command_atoms.exit => Some(exit_daemon(state, Some((conn, &ev)))),
                             _ => None,
                         };
 
@@ -579,23 +579,7 @@ fn handle_xcb_events<'a>(conn: &'a xcb::Connection, state: &mut State<'a>, comma
                             } else {
                                 true
                             };
-
+                            do_command_reply(conn, &ev, is_success);
-                            let reply_window = match ev.data() {
-                                x::ClientMessageData::Data32(data) => {
-                                    match unsafe { x::Window::new(data[0]) } {
-                                        x::WINDOW_NONE => None,
-                                        wid => Some(wid),
-                                    }
-                                },
-                                _ => None,
-                            };
-                            if let Some(reply_window) = reply_window {
-                                if let Err(err) = bscreensaver_command_response(conn, reply_window, is_success) {
-                                    info!("Failed to send command response: {}", err);
-                                }
-                            } else {
-                                debug!("Command sender did not include a reply window");
-                            }
                         }
                     },
                     xcb::Event::X(x::Event::MapNotify(ev)) if ev.window() == unlock_dialog_window(&state) => {
@@ -681,6 +665,25 @@ fn handle_xcb_events<'a>(conn: &'a xcb::Connection, state: &mut State<'a>, comma
     Ok(())
 }
 
+fn do_command_reply(conn: &xcb::Connection, ev: &x::ClientMessageEvent, is_success: bool) {
+    let reply_window = match ev.data() {
+        x::ClientMessageData::Data32(data) => {
+            match unsafe { x::Window::new(data[0]) } {
+                x::WINDOW_NONE => None,
+                wid => Some(wid),
+            }
+        },
+        _ => None,
+    };
+    if let Some(reply_window) = reply_window {
+        if let Err(err) = bscreensaver_command_response(conn, reply_window, is_success) {
+            info!("Failed to send command response: {}", err);
+        }
+    } else {
+        debug!("Command sender did not include a reply window");
+    }
+}
+
 fn embedder_window(state: &State) -> x::Window {
     state.unlock_dialog.as_ref().map(|ud| ud.embedder.embedder_window()).unwrap_or(x::WINDOW_NONE)
 }
@@ -940,7 +943,7 @@ fn unlock_screen(conn: &xcb::Connection, state: &mut State) -> anyhow::Result<()
     Ok(())
 }
 
-fn restart_daemon(state: &mut State) -> anyhow::Result<()> {
+fn restart_daemon(state: &mut State, trigger: Option<(&xcb::Connection, &x::ClientMessageEvent)>) -> anyhow::Result<()> {
     info!("Restarting");
 
     let exe = read_link("/proc/self/exe")?;
@@ -959,9 +962,17 @@ fn restart_daemon(state: &mut State) -> anyhow::Result<()> {
     match unsafe { fork() } {
         Err(err) => {
             error!("Failed to fork: {}", err);
+            if let Some((conn, ev)) = trigger {
+                do_command_reply(conn, ev, false);
+            }
             Err(err)?;
         },
-        Ok(ForkResult::Parent { .. }) => exit(0),
+        Ok(ForkResult::Parent { .. }) => {
+            if let Some((conn, ev)) = trigger {
+                do_command_reply(conn, ev, true);
+            }
+            exit(0);
+        },
         Ok(ForkResult::Child) => {
             if let Err(err) = setsid() {
                 warn!("Failed to start new session: {}", err);
@@ -973,11 +984,14 @@ fn restart_daemon(state: &mut State) -> anyhow::Result<()> {
     Ok(())
 }
 
-fn exit_daemon(state: &mut State) -> anyhow::Result<()> {
+fn exit_daemon(state: &mut State, trigger: Option<(&xcb::Connection, &x::ClientMessageEvent)>) -> anyhow::Result<()> {
     info!("Quitting");
     if let Err(err) = kill_child_processes(state) {
         warn!("Failed to kill child processes: {}", err);
     }
+    if let Some((conn, ev)) = trigger {
+        do_command_reply(conn, ev, true);
+    }
     exit(0);
 }