bscreensaver/locker/src/main.rs

use clap::{Arg, Command as ClapCommand};
use log::{debug, error, info, trace, warn};
use nix::{
    poll::{poll, PollFd, PollFlags},
    unistd::{execv, fork, setsid, ForkResult},
    sys::{
        signal::{sigprocmask, SigSet, SigmaskHow, Signal},
        signalfd::{SignalFd, SfdFlags},
    },
};
use std::{
    env,
    ffi::CString,
    fs::read_link,
    io,
    os::unix::io::AsRawFd,
    rc::Rc,
    process::exit,
    sync::Mutex,
    time::{Duration, Instant}, path::PathBuf,
};
use xcb::{randr, x, xfixes, xinput};

use bscreensaver::{screensaver::{BlankerState, CommandHandlers, Screensaver}, subservice::Subservices};
use bscreensaver_util::{*, settings::Configuration};

const BLANKED_ARG: &str = "blanked";
const LOCKED_ARG: &str = "locked";

macro_rules! maybe_add_fd {
    ($pfds:expr, $fd:expr) => {
        if let Some(fd) = $fd {
            $pfds.push(PollFd::new(fd, PollFlags::POLLIN));
            Some(fd)
        } else {
            None
        }
    };
}

fn main() -> anyhow::Result<()> {
    init_logging("BSCREENSAVER_LOG");

    let config = Configuration::load()?;

    let args = ClapCommand::new("Blanks and locks the screen after a period of time")
        .author(env!("CARGO_PKG_AUTHORS"))
        .version(env!("CARGO_PKG_VERSION"))
        .arg(
            Arg::new("blanked")
                .long(BLANKED_ARG)
                .help("Starts up in the blanked screensaver")
        )
        .arg(
            Arg::new("locked")
                .long(LOCKED_ARG)
                .help("Starts up in the blanked and locked screensaver")
        )
        .get_matches();

    let mut signal_fd = init_signals()?;

    let (conn, screen_num) = xcb::Connection::connect_with_extensions(
        None,
        &[xcb::Extension::RandR, xcb::Extension::XFixes, xcb::Extension::Input, xcb::Extension::Xkb],
        &[]
    )?;
    let setup = conn.get_setup();
    let screen = setup.roots().nth(screen_num as usize).unwrap();

    init_xfixes(&conn)?;
    init_xinput(&conn)?;
    init_randr(&conn)?;
    xkb::x11::setup(&conn, xkb::x11::MIN_MAJOR_XKB_VERSION, xkb::x11::MIN_MINOR_XKB_VERSION, xkb::x11::NO_FLAGS)
        .map_err(|_| anyhow::anyhow!("Failed to initialize XKB extension"))?;

    let helper_dir = PathBuf::from(env!("HELPER_DIR"));

    let subservices = Rc::new(Mutex::new(Subservices::start_all(&helper_dir)?));
    let command_handlers = {
        let subservices1 = Rc::clone(&subservices);
        let subservices2 = Rc::clone(&subservices);
        CommandHandlers {
            restart_handler: &move |screensaver, conn, trigger| {
                restart_daemon(screensaver, &mut subservices1.lock().unwrap(), trigger.map(|t| (conn, t)))
            },
            exit_handler: &move |screensaver, conn, trigger| {
                exit_daemon(screensaver, &mut subservices2.lock().unwrap(), trigger.map(|t| (conn, t)))
            },
        }
    };
    let mut screensaver = Screensaver::new(&config, &helper_dir, &command_handlers, &conn, screen)?;

    if args.is_present(LOCKED_ARG) {
        match screensaver.lock_screen(&conn) {
            Err(err) => error!("POSSIBLY FAILED TO LOCK SCREEN ON STARTUP: {}", err),
            Ok(_) => debug!("Got --{} arg; screen locked on startup", LOCKED_ARG),
        }
    } else if args.is_present(BLANKED_ARG) {
        match screensaver.blank_screen(&conn) {
            Err(err) => warn!("Possibly failed to blank screen on startup: {}", err),
            Ok(_) => debug!("Got --{} arg; screen locked on startup", BLANKED_ARG),
        }
    }

    let _ = conn.send_and_check_request(&x::SetScreenSaver {
        timeout: 0,
        interval: 0,
        prefer_blanking: x::Blanking::NotPreferred,
        allow_exposures: x::Exposures::NotAllowed,
    });


    loop {
        if let Err(err) = screensaver.handle_xcb_events(&conn) {
            if conn.has_error().is_err() {
                error!("Lost connection to X server; attempting to restart");
                (command_handlers.restart_handler)(&mut screensaver, &conn, None)?;
            }
            warn!("Error handling event: {}", err);
        }

        let mut pfds = Vec::new();
        pfds.push(PollFd::new(signal_fd.as_raw_fd(), PollFlags::POLLIN));
        pfds.push(PollFd::new(conn.as_raw_fd(), PollFlags::POLLIN));
        let dbus_service_fd = maybe_add_fd!(&mut pfds, subservices.lock().unwrap().dbus_service().map(|ds| ds.pidfd().as_raw_fd()));
        let systemd_service_fd = maybe_add_fd!(&mut pfds, subservices.lock().unwrap().systemd_service().map(|ds| ds.pidfd().as_raw_fd()));
        let dialog_fd = maybe_add_fd!(&mut pfds, screensaver.unlock_dialog_pidfd().map(|udpfd| udpfd.as_raw_fd()));

        let since_last_activity = Instant::now().duration_since(screensaver.last_user_activity());
        let poll_timeout = match screensaver.blanker_state() {
            BlankerState::Idle if since_last_activity > config.lock_timeout - config.blank_before_locking => Some(Duration::ZERO),
            BlankerState::Idle => Some(config.lock_timeout - config.blank_before_locking - since_last_activity),
            BlankerState::Blanked if since_last_activity > config.lock_timeout => Some(Duration::ZERO),
            BlankerState::Blanked => Some(config.lock_timeout - since_last_activity),
            BlankerState::Locked => None,
        };
        let poll_timeout = poll_timeout.map(|pt| if pt.as_millis() > i32::MAX as u128 {
            i32::MAX
        } else {
            pt.as_millis() as i32
        }).unwrap_or(-1);

        trace!("about to poll (timeout={})", poll_timeout);
        let nready = poll(pfds.as_mut_slice(), poll_timeout)?;  // FIXME: maybe shouldn't quit here on errors if screen is locked
        trace!("polled; {} FD ready", nready);
        if nready > 0 {
            for pfd in pfds {
                if pfd.revents().filter(|pf| pf.contains(PollFlags::POLLIN)).is_some() {
                    let result = match pfd.as_raw_fd() {
                        fd if fd == signal_fd.as_raw_fd() => handle_signals(&mut screensaver, &mut subservices.lock().unwrap(), &mut signal_fd),
                        fd if fd == conn.as_raw_fd() => screensaver.handle_xcb_events(&conn),
                        fd if opt_contains(&dbus_service_fd, &fd) => subservices.lock().unwrap().handle_quit(),
                        fd if opt_contains(&systemd_service_fd, &fd) => subservices.lock().unwrap().handle_quit(),
                        fd if opt_contains(&dialog_fd, &fd) => screensaver.handle_unlock_dialog_quit(&conn),
                        _ => Ok(()),
                    };

                    if let Err(err) = result {
                        if conn.has_error().is_err() {
                            error!("Lost connection to X server; atempting to restart");
                            (command_handlers.restart_handler)(&mut screensaver, &conn, None)?;
                        }
                        warn!("Error handling event: {}", err);
                    }
                }
            }
        }

        let since_last_activity = Instant::now().duration_since(screensaver.last_user_activity());

        if screensaver.blanker_state() < BlankerState::Blanked && since_last_activity > config.lock_timeout - config.blank_before_locking {
            if let Err(err) = screensaver.blank_screen(&conn) {
                error!("POSSIBLY FAILED TO BLANK SCREEN: {}", err);
            }
        }

        if screensaver.blanker_state() < BlankerState::Locked && since_last_activity > config.lock_timeout {
            if let Err(err) = screensaver.lock_screen(&conn) {
                error!("POSSIBLY FAILED TO LOCK SCREEN: {}", err);
            }
        }
    }
}

fn init_signals() -> anyhow::Result<SignalFd> {
    let sigs = {
        let mut s = SigSet::empty();
        s.add(Signal::SIGHUP);
        s.add(Signal::SIGINT);
        s.add(Signal::SIGQUIT);
        s.add(Signal::SIGTERM);
        s
    };
    sigprocmask(SigmaskHow::SIG_BLOCK, Some(&sigs), None)?;
    let flags = SfdFlags::SFD_NONBLOCK | SfdFlags::SFD_CLOEXEC;
    let fd = SignalFd::with_flags(&sigs, flags)?;
    Ok(fd)
}

fn init_xfixes(conn: &xcb::Connection) -> xcb::Result<()> {
    let cookie = conn.send_request(&xfixes::QueryVersion {
        client_major_version: 4,
        client_minor_version: 0,
    });
    let reply = conn.wait_for_reply(cookie)?;
    if reply.major_version() < 4 {
        warn!("XFIXES version 4 or better is not supported by the X server (max supported: {}.{})", reply.major_version(), reply.minor_version());
    }
    Ok(())
}

fn init_xinput(conn: &xcb::Connection) -> xcb::Result<()> {
    let cookie = conn.send_request(&xcb::xinput::XiQueryVersion {
        major_version: 2,
        minor_version: 2,
    });
    let reply = conn.wait_for_reply(cookie)?;
    if reply.major_version() < 2 {
        error!("Version 2 or greater of the Xinput extension is required (got {}.{})", reply.major_version(), reply.minor_version());
        exit(1);
    }

    let mut cookies = Vec::new();
    for screen in conn.get_setup().roots() {
        cookies.push(conn.send_request_checked(&xinput::XiSelectEvents {
            window: screen.root(),
            masks: &[
                xinput::EventMaskBuf::new(
                    xinput::Device::AllMaster,
                    &[
                        xinput::XiEventMask::RAW_KEY_PRESS |
                        xinput::XiEventMask::RAW_KEY_RELEASE |
                        xinput::XiEventMask::RAW_BUTTON_PRESS |
                        xinput::XiEventMask::RAW_BUTTON_RELEASE |
                        xinput::XiEventMask::RAW_TOUCH_BEGIN |
                        xinput::XiEventMask::RAW_TOUCH_UPDATE |
                        xinput::XiEventMask::RAW_TOUCH_END |
                        xinput::XiEventMask::RAW_MOTION
                    ]
                )
            ],
        }));
    }
    for cookie in cookies {
        conn.check_request(cookie)?;
    }

    Ok(())
}

fn init_randr(conn: &xcb::Connection) -> anyhow::Result<()> {
    let cookie = conn.send_request(&randr::QueryVersion {
        major_version: 1,
        minor_version: 2,
    });
    let reply = conn.wait_for_reply(cookie)?;
    if reply.major_version() < 1 || (reply.major_version() == 1 && reply.minor_version() < 2) {
        Err(anyhow::anyhow!("XRandR extension version 1.2 or greater is required, but the X server only supports {}.{}", reply.major_version(), reply.minor_version()))?;
    }

    for screen in conn.get_setup().roots() {
        conn.send_and_check_request(&randr::SelectInput {
            window: screen.root(),
            enable: randr::NotifyMask::SCREEN_CHANGE | randr::NotifyMask::CRTC_CHANGE | randr::NotifyMask::OUTPUT_CHANGE,
        })?;
    }
    Ok(())
}

fn handle_signals(screensaver: &mut Screensaver, subservices: &mut Subservices, signal_fd: &mut SignalFd) -> anyhow::Result<()> {
    match signal_fd.read_signal()? {
        None => (),
        Some(info) if info.ssi_signo == Signal::SIGHUP as u32 => restart_daemon(screensaver, subservices, None)?,
        Some(info) if info.ssi_signo == Signal::SIGINT as u32 => exit_daemon(screensaver, subservices, None)?,
        Some(info) if info.ssi_signo == Signal::SIGQUIT as u32 => exit_daemon(screensaver, subservices, None)?,
        Some(info) if info.ssi_signo == Signal::SIGTERM as u32 => exit_daemon(screensaver, subservices, None)?,
        Some(info) => trace!("Unexpected signal {}", info.ssi_signo),
    }
    Ok(())
}

fn restart_daemon(screensaver: &mut Screensaver, subservices: &mut Subservices, trigger: Option<(&xcb::Connection, &x::ClientMessageEvent)>) -> anyhow::Result<()> {
    info!("Restarting");

    let exe = read_link("/proc/self/exe")?;
    let exe = CString::new(exe.to_str().ok_or(io::Error::new(io::ErrorKind::InvalidData, "Path cannot be converted to str".to_string()))?)?;

    let argv = match screensaver.blanker_state() {
        BlankerState::Idle => vec![],
        BlankerState::Blanked => vec![CString::new("--".to_string() + BLANKED_ARG)?],
        BlankerState::Locked => vec![CString::new("--".to_string() + LOCKED_ARG)?],
    };

    kill_child_processes(screensaver, subservices);

    match unsafe { fork() } {
        Err(err) => {
            error!("Failed to fork: {}", err);
            if let Some((conn, ev)) = trigger {
                bscreensaver::send_command_reply(conn, ev, false);
            }
            Err(err)?;
        },
        Ok(ForkResult::Parent { .. }) => {
            if let Some((conn, ev)) = trigger {
                bscreensaver::send_command_reply(conn, ev, true);
            }
            exit(0);
        },
        Ok(ForkResult::Child) => {
            if let Err(err) = setsid() {
                warn!("Failed to start new session: {}", err);
            }
            execv(exe.as_c_str(), &argv)?;
        },
    }

    Ok(())
}

fn exit_daemon(screensaver: &mut Screensaver, subservices: &mut Subservices, trigger: Option<(&xcb::Connection, &x::ClientMessageEvent)>) -> anyhow::Result<()> {
    info!("Quitting");
    kill_child_processes(screensaver, subservices);
    if let Some((conn, ev)) = trigger {
        bscreensaver::send_command_reply(conn, ev, true);
    }
    exit(0);
}

fn kill_child_processes(screensaver: &mut Screensaver, subservices: &mut Subservices) {
    screensaver.stop_unlock_dialog();
    subservices.stop_all();
}