gtk-ahfail/crates/ahfail-pam/src/lib.rs

use libc::{c_char, c_int, c_void, pid_t};
use std::ffi::CString;

// --- PAM constants (platform-specific) ---

pub const PAM_SUCCESS: c_int = 0;
pub const PAM_IGNORE: c_int = 25;

#[cfg(target_os = "linux")]
pub const PAM_AUTH_ERR: c_int = 7;
#[cfg(target_os = "macos")]
pub const PAM_AUTH_ERR: c_int = 9;
#[cfg(not(any(target_os = "linux", target_os = "macos")))]
pub const PAM_AUTH_ERR: c_int = 7;

#[cfg(target_os = "linux")]
pub const PAM_DATA_REPLACE: c_int = 0x20000000u32 as i32;
#[cfg(target_os = "macos")]
pub const PAM_DATA_REPLACE: c_int = 0x00000002;
#[cfg(not(any(target_os = "linux", target_os = "macos")))]
pub const PAM_DATA_REPLACE: c_int = 0x20000000u32 as i32;

// Default install path for ahfail-display.
// On Linux: built from AHFAIL_INSTALL_DIR emitted by build.rs, which reads AHFAIL_LIBDIR
// passed by Meson — correct on multiarch systems (e.g. /usr/lib/x86_64-linux-gnu/ahfail).
#[cfg(target_os = "macos")]
const DEFAULT_PATH: &str = "/usr/local/lib/ahfail/ahfail-display";
#[cfg(not(target_os = "macos"))]
const DEFAULT_PATH: &str = concat!(env!("AHFAIL_INSTALL_DIR"), "/ahfail-display");

pub fn default_display_path() -> &'static str { DEFAULT_PATH }

pub fn is_failure(s: c_int) -> bool { s & !PAM_DATA_REPLACE == PAM_AUTH_ERR }
pub fn is_success(s: c_int) -> bool { s & !PAM_DATA_REPLACE == PAM_SUCCESS }
pub fn is_replace(s: c_int) -> bool { s & PAM_DATA_REPLACE != 0 }

// --- PAM opaque handle ---

#[repr(C)]
pub struct PamHandle { _private: [u8; 0] }

type CleanupFn = unsafe extern "C" fn(*mut PamHandle, *mut c_void, c_int);

extern "C" {
    fn pam_set_data(
        pamh: *mut PamHandle,
        name: *const c_char,
        data: *mut c_void,
        cleanup: Option<CleanupFn>,
    ) -> c_int;
}

// --- Cleanup: fires on pam_end() or when data is replaced ---
unsafe extern "C" fn ahfail_cleanup(
    _pamh: *mut PamHandle,
    data: *mut c_void,
    error_status: c_int,
) {
    // Reclaim any stored path string to avoid a leak (PAM guarantees this fires exactly once).
    let path_override: Option<String> = if data.is_null() {
        None
    } else {
        Some(*Box::from_raw(data as *mut String))
    };

    if is_replace(error_status) {
        // Data replaced — a new pam_set_data call overwrote ours. Only spawn on failure;
        // on success (PAM_SUCCESS | PAM_DATA_REPLACE) we do nothing.
        if is_failure(error_status) {
            spawn_display(path_override);
        }
        return;
    }
    if is_failure(error_status) {
        spawn_display(path_override);
    } else if is_success(error_status) {
        kill_display();
    }
}

// --- Exported PAM module entry points ---

#[no_mangle]
pub unsafe extern "C" fn pam_sm_authenticate(
    pamh: *mut PamHandle,
    _flags: c_int,
    argc: c_int,
    argv: *const *const c_char,
) -> c_int {
    if pamh.is_null() { return PAM_IGNORE; }

    let args = argc_argv(argc, argv);
    let display_path = read_display_path_arg(args);

    let key = match CString::new("dk.weircon.ahfail") {
        Ok(k) => k,
        Err(_) => return PAM_IGNORE,
    };

    // Store optional path override as heap data for the cleanup function.
    let path_ptr: *mut c_void = match display_path {
        Some(p) => Box::into_raw(Box::new(p)) as *mut c_void,
        None => std::ptr::null_mut(),
    };

    let ret = pam_set_data(pamh, key.as_ptr(), path_ptr, Some(ahfail_cleanup));
    if ret != PAM_SUCCESS && !path_ptr.is_null() {
        // pam_set_data failed — cleanup will never fire, so free the Box ourselves.
        drop(Box::from_raw(path_ptr as *mut String));
    }
    PAM_IGNORE
}

#[no_mangle]
pub unsafe extern "C" fn pam_sm_setcred(
    _pamh: *mut PamHandle, _flags: c_int, _argc: c_int, _argv: *const *const c_char,
) -> c_int { PAM_IGNORE }

#[no_mangle]
pub unsafe extern "C" fn pam_sm_acct_mgmt(
    _pamh: *mut PamHandle, _flags: c_int, _argc: c_int, _argv: *const *const c_char,
) -> c_int { PAM_IGNORE }

#[no_mangle]
pub unsafe extern "C" fn pam_sm_open_session(
    _pamh: *mut PamHandle, _flags: c_int, _argc: c_int, _argv: *const *const c_char,
) -> c_int { PAM_IGNORE }

#[no_mangle]
pub unsafe extern "C" fn pam_sm_close_session(
    _pamh: *mut PamHandle, _flags: c_int, _argc: c_int, _argv: *const *const c_char,
) -> c_int { PAM_IGNORE }

#[no_mangle]
pub unsafe extern "C" fn pam_sm_chauthtok(
    _pamh: *mut PamHandle, _flags: c_int, _argc: c_int, _argv: *const *const c_char,
) -> c_int { PAM_IGNORE }

// --- Helpers ---

unsafe fn argc_argv<'a>(argc: c_int, argv: *const *const c_char) -> &'a [*const c_char] {
    if argv.is_null() || argc <= 0 { return &[]; }
    std::slice::from_raw_parts(argv, argc as usize)
}

fn read_display_path_arg(args: &[*const c_char]) -> Option<String> {
    let prefix = b"display_path=";
    for &arg in args {
        if arg.is_null() { continue; }
        let s = unsafe { std::ffi::CStr::from_ptr(arg) }.to_bytes();
        if s.starts_with(prefix) {
            return std::str::from_utf8(&s[prefix.len()..]).ok().map(|s| s.to_string());
        }
    }
    None
}

fn spawn_display(path_override: Option<String>) {
    let path = path_override.unwrap_or_else(|| DEFAULT_PATH.to_string());
    let cpath = match CString::new(path.as_str()) {
        Ok(p) => p,
        Err(_) => return,
    };
    // Double-fork: grandchild adopted by init, PAM stack not blocked.
    unsafe {
        let pid: pid_t = libc::fork();
        if pid < 0 { return; }
        if pid > 0 {
            // Parent: wait for intermediate to exit immediately.
            libc::waitpid(pid, std::ptr::null_mut(), 0);
            return;
        }
        // Intermediate: fork again then exit immediately.
        let pid2: pid_t = libc::fork();
        if pid2 != 0 { libc::_exit(0); }
        // Grandchild: detach from PAM daemon's session, close all inherited fds, exec.
        libc::setsid();
        let max_fd = libc::sysconf(libc::_SC_OPEN_MAX) as c_int;
        for fd in 0..max_fd.min(4096) {
            libc::close(fd);
        }
        let args: [*const c_char; 2] = [cpath.as_ptr(), std::ptr::null()];
        libc::execv(cpath.as_ptr(), args.as_ptr());
        libc::_exit(1);
    }
}

fn kill_display() {
    // SIGTERM all ahfail-display processes owned by the current user.
    let uid = unsafe { libc::getuid() }.to_string();
    let _ = std::process::Command::new("pkill")
        .args(["-SIGTERM", "-u", &uid, "ahfail-display"])
        .spawn();
}