// Copyright (C) 2023  Tomasz Kramkowski <tomasz@kramkow.ski>
// SPDX-License-Identifier: MIT

mod pam;
mod passwd;
mod syslog;

use pam::PAMHandle;
use rustix::fd::{AsFd, OwnedFd};
use rustix::fs::{self, Mode, OFlags};
use rustix::io::{self, Errno};
use rustix::path::Arg;
use std::ffi::{c_char, c_int, CStr, CString};
use std::panic;
use std::process;
use syslog::{Facility, Level, Priority};

fn create_and_open_dir<D: AsFd, P: Arg + Copy>(
    d: &D,
    path: P,
) -> rustix::io::Result<OwnedFd> {
    if let Err(e) = fs::mkdirat(d, path, Mode::RWXU | Mode::RWXG | Mode::RWXO) {
        if e != Errno::EXIST {
            return Err(e);
        }
    }
    fs::openat(d, path, OFlags::DIRECTORY | OFlags::PATH, Mode::empty())
}

struct SessionError;

impl From<rustix::io::Errno> for SessionError {
    fn from(_: rustix::io::Errno) -> Self {
        SessionError
    }
}

impl From<std::ffi::NulError> for SessionError {
    fn from(_: std::ffi::NulError) -> Self {
        SessionError
    }
}

const PRIORITY: Priority = Priority {
    level: Level::Debug,
    facility: Facility::Auth,
};

fn open_session(h: &PAMHandle, mountpoint: &str) -> Result<(), SessionError> {
    let user = match h.get_user::<CStr>(None) {
        Ok(user) => user,
        Err(e) => {
            if let Ok(message) =
                CString::new(format!("Failure to get username: {e}"))
            {
                h.syslog(PRIORITY, &message);
            }
            return Err(SessionError);
        }
    };
    let uid = match passwd::get_uid_by_name(&user) {
        Some(uid) => CString::new(uid.to_string())?,
        None => {
            if let Ok(message) =
                CString::new(format!("Failure to map user {user:?} to passwd entry"))
            {
                h.syslog(PRIORITY, &message);
            }
            return Err(SessionError);
        }
    };
    let d = fs::openat(
        fs::cwd(),
        mountpoint,
        OFlags::DIRECTORY | OFlags::PATH,
        Mode::empty(),
    )?;
    let d = create_and_open_dir(&d, "user")?;
    let d = create_and_open_dir(&d, &uid)?;
    let d = create_and_open_dir(&d, "leaf")?;
    let pid = CString::new(process::id().to_string())?;
    let pid = pid.as_bytes();
    let procs = fs::openat(d, "cgroup.procs", OFlags::WRONLY, Mode::empty())?;
    match io::write(procs, pid) {
        Ok(n) if n == pid.len() => Ok(()),
        _ => Err(SessionError),
    }
}

const CG_MOUNT: &str = "/sys/fs/cgroup";

#[no_mangle]
pub extern "C" fn pam_sm_open_session(
    h: &mut PAMHandle,
    _flags: c_int,
    _argc: c_int,
    _argv: *const *const c_char,
) -> c_int {
    match panic::catch_unwind(|| open_session(h, CG_MOUNT)) {
        Ok(Ok(())) => pam::SUCCESS,
        _ => pam::SESSION_ERR,
    }
}

#[no_mangle]
pub extern "C" fn pam_sm_close_session(
    _h: &mut PAMHandle,
    _flags: c_int,
    _argc: c_int,
    _argv: *const *const c_char,
) -> c_int {
    pam::IGNORE
}