2 files changed, 69 insertions, 32 deletions

diff --git a/src/bin/agent.rs b/src/bin/agent.rs
index 50b13a1..f21abcc 100644
--- a/src/bin/agent.rs
+++ b/src/bin/agent.rs
@@ -1,35 +1,15 @@
-use fs2::FileExt as _;
-use std::io::Write as _;
 use tokio::io::AsyncBufReadExt as _;
 use tokio::stream::StreamExt as _;
 
-fn make_pidfile() -> std::fs::File {
+fn make_socket() -> anyhow::Result<tokio::net::UnixListener> {
     let runtime_dir = rbw::dirs::runtime_dir();
-    std::fs::create_dir_all(&runtime_dir).unwrap();
-
-    let mut fh = std::fs::OpenOptions::new()
-        .write(true)
-        .create(true)
-        .open(runtime_dir.join("pidfile"))
-        .unwrap();
-
-    fh.try_lock_exclusive().unwrap();
-
-    fh.set_len(0).unwrap();
-    fh.write_all(format!("{}", std::process::id()).as_bytes())
-        .unwrap();
-
-    fh
-}
-
-fn make_socket() -> tokio::net::UnixListener {
-    let runtime_dir = rbw::dirs::runtime_dir();
-    std::fs::create_dir_all(&runtime_dir).unwrap();
+    std::fs::create_dir_all(&runtime_dir)?;
 
     let path = runtime_dir.join("socket");
-    std::fs::remove_file(&path).unwrap();
-
-    tokio::net::UnixListener::bind(path).unwrap()
+    std::fs::remove_file(&path)?;
+    let sock = tokio::net::UnixListener::bind(&path)?;
+    log::debug!("listening on socket {}", path.to_string_lossy());
+    Ok(sock)
 }
 
 async fn ensure_login(state: std::sync::Arc<tokio::sync::RwLock<State>>) {
@@ -181,10 +161,49 @@ impl Agent {
     }
 }
 
-#[tokio::main]
-async fn main() {
-    let _lock = make_pidfile();
-    let listener = make_socket();
-    let mut agent = Agent::new();
-    agent.run(listener).await;
+fn main() {
+    env_logger::from_env(
+        env_logger::Env::default().default_filter_or("info"),
+    )
+    .init();
+
+    let runtime_dir = rbw::dirs::runtime_dir();
+    std::fs::create_dir_all(&runtime_dir).unwrap();
+
+    let (r, w) = nix::unistd::pipe().unwrap();
+    let res = daemonize::Daemonize::new()
+        .pid_file(runtime_dir.join("pidfile"))
+        .exit_action(move || {
+            nix::unistd::close(w).unwrap();
+            let mut buf = [0; 1];
+            nix::unistd::read(r, &mut buf).unwrap();
+            nix::unistd::close(r).unwrap();
+        })
+        .start();
+    nix::unistd::close(r).unwrap();
+
+    match res {
+        Ok(_) => (),
+        Err(e) => {
+            match e {
+                daemonize::DaemonizeError::LockPidfile(_) => {
+                    // this means that there is already an agent running, so
+                    // return a special exit code to allow the cli to detect
+                    // this case and not error out
+                    std::process::exit(23);
+                }
+                _ => panic!("failed to daemonize: {}", e),
+            }
+        }
+    }
+
+    tokio::runtime::Runtime::new().unwrap().block_on(async {
+        let listener = make_socket();
+
+        nix::unistd::write(w, &[0]).unwrap();
+        nix::unistd::close(w).unwrap();
+
+        let mut agent = Agent::new();
+        agent.run(listener.unwrap()).await;
+    })
 }
diff --git a/src/bin/rbw.rs b/src/bin/rbw.rs
index 910350b..a110176 100644
--- a/src/bin/rbw.rs
+++ b/src/bin/rbw.rs
@@ -1,5 +1,21 @@
 use std::io::Write as _;
 
+fn ensure_agent() {
+    let agent_path = std::env::var("RBW_AGENT");
+    let agent_path = agent_path
+        .as_ref()
+        .map(|s| s.as_str())
+        .unwrap_or("rbw-agent");
+    let status = std::process::Command::new(agent_path).status().unwrap();
+    if !status.success() {
+        if let Some(code) = status.code() {
+            if code != 23 {
+                panic!("failed to run agent: {}", status);
+            }
+        }
+    }
+}
+
 fn send(msg: &rbw::agent::Message) {
     let mut sock = std::os::unix::net::UnixStream::connect(
         rbw::dirs::runtime_dir().join("socket"),
@@ -80,6 +96,8 @@ fn main() {
         .subcommand(clap::SubCommand::with_name("purge"))
         .get_matches();
 
+    ensure_agent();
+
     match matches.subcommand() {
         ("login", Some(_)) => login(),
         ("unlock", Some(_)) => unlock(),