use std::{
env,
ffi::c_void,
fmt,
fs::{self, File},
io::Write,
path::{Path, PathBuf},
process::Command,
ptr,
sync::{
atomic::{AtomicBool, Ordering},
mpsc, Arc, Mutex,
},
thread,
};
use chrono;
use ctrlc;
use log::info;
use timer;
use core_foundation::{
array::{kCFTypeArrayCallBacks, CFArray, CFArrayCreate, CFArrayRef},
base::{CFAllocatorRef, CFRelease, CFType, CFTypeRef, TCFType, ToVoid},
runloop::{
kCFRunLoopDefaultMode, CFRunLoopAddSource, CFRunLoopGetCurrent, CFRunLoopRef, CFRunLoopRun,
CFRunLoopStop,
},
string::{CFString, CFStringRef},
};
use system_configuration_sys::{
dynamic_store::{
SCDynamicStoreContext, SCDynamicStoreCreate, SCDynamicStoreCreateRunLoopSource,
SCDynamicStoreRef, SCDynamicStoreSetNotificationKeys,
},
dynamic_store_copy_specific::{uid_t, SCDynamicStoreCopyConsoleUser},
};
use crate::controller::{ControllerInterface, ServiceMainFn};
use crate::session;
use crate::Error;
use crate::ServiceEvent;
type MacosServiceMainWrapperFn = extern "system" fn(args: Vec<String>);
pub type Session = session::Session_<u32>;
pub enum LaunchAgentTargetSesssion {
GUI,
NonGUI,
PerUser,
PreLogin,
}
impl fmt::Display for LaunchAgentTargetSesssion {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
match self {
LaunchAgentTargetSesssion::GUI => write!(f, "Aqua"),
LaunchAgentTargetSesssion::NonGUI => write!(f, "StandardIO"),
LaunchAgentTargetSesssion::PerUser => write!(f, "Background"),
LaunchAgentTargetSesssion::PreLogin => write!(f, "LoginWindow"),
}
}
}
fn launchctl_load_daemon(plist_path: &Path) -> Result<(), Error> {
let output = Command::new("launchctl")
.arg("load")
.arg(&plist_path.to_str().unwrap())
.output()
.map_err(|e| {
Error::new(&format!(
"Failed to load plist {}: {}",
plist_path.display(),
e
))
})?;
if output.stdout.len() > 0 {
info!("{}", String::from_utf8_lossy(&output.stdout));
}
Ok(())
}
fn launchctl_unload_daemon(plist_path: &Path) -> Result<(), Error> {
let output = Command::new("launchctl")
.arg("unload")
.arg(&plist_path.to_str().unwrap())
.output()
.map_err(|e| {
Error::new(&format!(
"Failed to unload plist {}: {}",
plist_path.display(),
e
))
})?;
if output.stdout.len() > 0 {
info!("{}", String::from_utf8_lossy(&output.stdout));
}
Ok(())
}
fn launchctl_start_daemon(name: &str) -> Result<(), Error> {
let output = Command::new("launchctl")
.arg("start")
.arg(name)
.output()
.map_err(|e| Error::new(&format!("Failed to start {}: {}", name, e)))?;
if output.stdout.len() > 0 {
info!("{}", String::from_utf8_lossy(&output.stdout));
}
Ok(())
}
fn launchctl_stop_daemon(name: &str) -> Result<(), Error> {
let output = Command::new("launchctl")
.arg("stop")
.arg(name)
.output()
.map_err(|e| Error::new(&format!("Failed to stop {}: {}", name, e)))?;
if output.stdout.len() > 0 {
info!("{}", String::from_utf8_lossy(&output.stdout));
}
Ok(())
}
pub struct MacosController {
pub service_name: String,
pub display_name: String,
pub description: String,
pub is_agent: bool,
pub session_types: Option<Vec<LaunchAgentTargetSesssion>>,
pub keep_alive: bool,
}
impl MacosController {
pub fn new(service_name: &str, display_name: &str, description: &str) -> MacosController {
MacosController {
service_name: service_name.to_string(),
display_name: display_name.to_string(),
description: description.to_string(),
is_agent: false,
session_types: None,
keep_alive: true,
}
}
pub fn register(
&mut self,
service_main_wrapper: MacosServiceMainWrapperFn,
) -> Result<(), Error> {
service_main_wrapper(env::args().collect());
Ok(())
}
fn get_plist_content(&self) -> Result<String, Error> {
let mut current_exe = env::current_exe()
.map_err(|e| Error::new(&format!("env::current_exe() failed: {}", e)))?;
let current_exe_str = current_exe
.to_str()
.expect("current_exe path to be unicode")
.to_string();
current_exe.pop();
let working_dir_str = current_exe
.to_str()
.expect("working_dir path to be unicode");
let mut plist = String::new();
plist.push_str(r#"
<?xml version="1.0" encoding="UTF-8"?>
<!DOCTYPE plist PUBLIC "-//Apple//DTD PLIST 1.0//EN" "http://www.apple.com/DTDs/PropertyList-1.0.dtd">
<plist version="1.0">
<dict>"#);
plist.push_str(&format!(
r#"
<key>Disabled</key>
<false/>
<key>Label</key>
<string>{}</string>
<key>ProgramArguments</key>
<array>
<string>{}</string>
</array>
<key>WorkingDirectory</key>
<string>{}</string>
<key>RunAtLoad</key>
<true/>"#,
self.service_name, current_exe_str, working_dir_str,
));
if self.is_agent {
if let Some(session_types) = self.session_types.as_ref() {
plist.push_str(
r#"
<key>LimitLoadToSessionType</key>
<array>"#,
);
for session_type in session_types {
plist.push_str(&format!(
r#"
<string>{}</string>"#,
session_type
));
}
plist.push_str(
r#"
</array>"#,
);
}
}
if self.keep_alive {
plist.push_str(
r#"
<key>KeepAlive</key>
<true/>"#,
);
}
plist.push_str(
r#"
</dict>
</plist>"#,
);
Ok(plist)
}
fn write_plist(&self, path: &Path) -> Result<(), Error> {
info!("Writing plist file {}", path.display());
let content = self.get_plist_content()?;
File::create(path)
.and_then(|mut file| file.write_all(content.as_bytes()))
.map_err(|e| Error::new(&format!("Failed to write {}: {}", path.display(), e)))
}
fn plist_path(&mut self) -> PathBuf {
Path::new("/Library/")
.join(if self.is_agent {
"LaunchAgents/"
} else {
"LaunchDaemons/"
})
.join(format!("{}.plist", &self.service_name))
}
}
impl ControllerInterface for MacosController {
fn create(&mut self) -> Result<(), Error> {
let plist_path = self.plist_path();
self.write_plist(&plist_path)?;
if !self.is_agent {
return launchctl_load_daemon(&plist_path);
}
Ok(())
}
fn delete(&mut self) -> Result<(), Error> {
let plist_path = self.plist_path();
if !self.is_agent {
launchctl_unload_daemon(&plist_path)?;
}
fs::remove_file(&plist_path)
.map_err(|e| Error::new(&format!("Failed to delete {}: {}", plist_path.display(), e)))
}
fn start(&mut self) -> Result<(), Error> {
launchctl_start_daemon(&self.service_name)
}
fn stop(&mut self) -> Result<(), Error> {
launchctl_stop_daemon(&self.service_name)
}
fn load(&mut self) -> Result<(), Error> {
launchctl_load_daemon(&self.plist_path())
}
fn unload(&mut self) -> Result<(), Error> {
launchctl_unload_daemon(&self.plist_path())
}
}
#[macro_export]
macro_rules! Service {
($name:expr, $function:ident) => {
extern "system" fn service_main_wrapper(args: Vec<String>) {
dispatch($function, args);
}
};
}
fn active_session_uid(store_ref: Option<SCDynamicStoreRef>) -> u32 {
let mut uid: uid_t = 0;
let store = store_ref.unwrap_or(ptr::null());
let user = unsafe { SCDynamicStoreCopyConsoleUser(store, &mut uid, ptr::null_mut()) };
if user.is_null() {
return 0;
}
let _cf_user: CFString = unsafe { TCFType::wrap_under_create_rule(user) };
return uid;
}
unsafe extern "C" fn on_sc_console_user_change<F>(
store: SCDynamicStoreRef,
_keys: CFArrayRef,
info: *mut c_void,
) where
F: FnMut(u32, EventType) + Send + 'static,
{
let uid = active_session_uid(Some(store));
let ctx_box = Box::from_raw(info as *mut Arc<SyncSessionContext<F>>);
let ctx_ptr = Box::leak(ctx_box);
let mut ctx = ctx_ptr.lock().unwrap();
let old_uid = ctx.uid;
if uid != ctx.uid {
ctx.uid = uid;
if old_uid != 0 || ctx.last_was_logout.load(Ordering::SeqCst) {
ctx.last_was_logout.store(false, Ordering::SeqCst);
(ctx.callback)(old_uid, EventType::Disconnect);
}
if uid != 0 {
ctx.pending_connect.store(false, Ordering::SeqCst);
(ctx.callback)(uid, EventType::Connect);
} else {
ctx.pending_connect.store(true, Ordering::SeqCst);
let ctx_weak = Arc::downgrade(ctx_ptr);
thread::spawn(move || {
let timer = timer::Timer::new();
let (tx, rx) = mpsc::channel();
let _guard = timer.schedule_with_delay(chrono::Duration::seconds(3), move || {
let _ignored = tx.send(());
});
rx.recv().unwrap();
match ctx_weak.upgrade() {
Some(ctx_ptr) => {
let mut ctx = ctx_ptr.lock().unwrap();
if ctx.pending_connect.load(Ordering::SeqCst) {
ctx.last_was_logout.store(true, Ordering::SeqCst);
(ctx.callback)(uid, EventType::Connect);
}
}
None => return,
};
});
}
}
}
#[link(name = "SystemConfiguration", kind = "framework")]
extern "C" {
pub fn SCDynamicStoreKeyCreateConsoleUser(allocator: CFAllocatorRef) -> CFStringRef;
}
pub enum EventType {
Connect,
Disconnect,
}
pub struct SessionContext<F: FnMut(u32, EventType)> {
uid: u32,
callback: F,
pending_connect: AtomicBool,
last_was_logout: AtomicBool,
}
impl<F: FnMut(u32, EventType)> SessionContext<F> {
pub fn new(cb: F) -> Self
where
F: FnMut(u32, EventType) + Send + 'static,
{
Self {
uid: active_session_uid(None),
callback: cb,
pending_connect: AtomicBool::new(false),
last_was_logout: AtomicBool::new(true),
}
}
}
pub type SyncSessionContext<T> = Mutex<SessionContext<T>>;
pub struct Monitor<F: FnMut(u32, EventType)> {
context: *mut Arc<SyncSessionContext<F>>,
}
impl<F: FnMut(u32, EventType)> Drop for Monitor<F> {
fn drop(&mut self) {
let _ = unsafe { Box::from_raw(self.context as *mut Arc<SyncSessionContext<F>>) };
}
}
impl<F: FnMut(u32, EventType)> Monitor<F> {
pub fn new(cb: F) -> Result<Self, std::io::Error>
where
F: FnMut(u32, EventType) + Send + 'static,
{
let session_ctx = Box::new(Arc::new(Mutex::new(SessionContext::new(cb))));
let session_ctx_ptr = Box::into_raw(session_ctx);
let mut ctx = SCDynamicStoreContext {
version: 0,
info: session_ctx_ptr as *mut c_void,
retain: None,
release: None,
copyDescription: None,
};
unsafe {
let name = CFString::from_static_string("kCGSSessionUserNameKey");
let store_ref = SCDynamicStoreCreate(
ptr::null_mut(),
name.to_void() as CFStringRef,
Some(on_sc_console_user_change::<F>),
&mut ctx,
);
let key = SCDynamicStoreKeyCreateConsoleUser(ptr::null());
let keys = CFArrayCreate(ptr::null(), &key.to_void(), 1, &kCFTypeArrayCallBacks);
let _ = SCDynamicStoreSetNotificationKeys(store_ref, keys, ptr::null_mut());
let _: CFArray<CFType> = TCFType::wrap_under_create_rule(keys);
let rls = SCDynamicStoreCreateRunLoopSource(ptr::null_mut(), store_ref, 0);
CFRunLoopAddSource(CFRunLoopGetCurrent(), rls, kCFRunLoopDefaultMode);
CFRelease(rls as CFTypeRef);
CFRelease(store_ref as CFTypeRef);
Ok(Monitor {
context: session_ctx_ptr,
})
}
}
}
pub struct MonitorLoopRef {
loop_ref: CFRunLoopRef,
}
unsafe impl Send for MonitorLoopRef {}
impl MonitorLoopRef {
pub fn stop(&mut self) {
unsafe { CFRunLoopStop(self.loop_ref) };
}
}
pub fn run_monitor<T: Send + 'static>(
tx: mpsc::Sender<ServiceEvent<T>>,
) -> Result<MonitorLoopRef, std::io::Error> {
let (_tx, rx) = mpsc::channel();
thread::spawn(move || {
let mon = Monitor::new(move |uid: u32, event: EventType| {
match event {
EventType::Connect => {
let _ = tx.send(ServiceEvent::SessionConnect(Session::new(uid)));
}
EventType::Disconnect => {
let _ = tx.send(ServiceEvent::SessionDisconnect(Session::new(uid)));
}
};
});
let loop_ref = unsafe { CFRunLoopGetCurrent() };
let mon_loop_ref = MonitorLoopRef { loop_ref: loop_ref };
_tx.send(mon_loop_ref).unwrap();
unsafe { CFRunLoopRun() };
drop(mon);
});
let received = rx.recv().unwrap();
return Ok(received);
}
#[doc(hidden)]
pub fn dispatch<T: Send + 'static>(service_main: ServiceMainFn<T>, args: Vec<String>) {
let (tx, rx) = mpsc::channel();
let mut session_monitor = run_monitor(tx.clone()).expect("Failed to run session monitor");
let _tx = tx.clone();
ctrlc::set_handler(move || {
let _ = tx.send(ServiceEvent::Stop);
})
.expect("Failed to register Ctrl-C handler");
service_main(rx, _tx, args, false);
session_monitor.stop();
}