From 8ce28d2d5326cdc15328cc3a88d652a758563ece Mon Sep 17 00:00:00 2001 From: Jesse Luehrs Date: Sat, 11 Apr 2020 19:32:39 -0400 Subject: also refactor the agent code --- src/bin/rbw-agent.rs | 248 ------------------------------------------- src/bin/rbw-agent/actions.rs | 116 ++++++++++++++++++++ src/bin/rbw-agent/agent.rs | 74 +++++++++++++ src/bin/rbw-agent/daemon.rs | 50 +++++++++ src/bin/rbw-agent/main.rs | 24 +++++ src/bin/rbw-agent/sock.rs | 36 +++++++ 6 files changed, 300 insertions(+), 248 deletions(-) delete mode 100644 src/bin/rbw-agent.rs create mode 100644 src/bin/rbw-agent/actions.rs create mode 100644 src/bin/rbw-agent/agent.rs create mode 100644 src/bin/rbw-agent/daemon.rs create mode 100644 src/bin/rbw-agent/main.rs create mode 100644 src/bin/rbw-agent/sock.rs diff --git a/src/bin/rbw-agent.rs b/src/bin/rbw-agent.rs deleted file mode 100644 index 00c4dd5..0000000 --- a/src/bin/rbw-agent.rs +++ /dev/null @@ -1,248 +0,0 @@ -use tokio::io::{AsyncBufReadExt as _, AsyncWriteExt as _}; -use tokio::stream::StreamExt as _; - -fn make_socket() -> anyhow::Result { - let runtime_dir = rbw::dirs::runtime_dir(); - std::fs::create_dir_all(&runtime_dir)?; - - let path = runtime_dir.join("socket"); - 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 send_response( - sock: &mut tokio::net::UnixStream, - res: &rbw::agent::Response, -) { - sock.write_all(serde_json::to_string(res).unwrap().as_bytes()) - .await - .unwrap(); - sock.write_all(b"\n").await.unwrap(); -} - -async fn login( - sock: &mut tokio::net::UnixStream, - state: std::sync::Arc>, - tty: Option<&str>, -) { - let mut state = state.write().await; - let email = config_email().await; - let mut db = rbw::db::Db::load_async(&email) - .await - .unwrap_or_else(|_| rbw::db::Db::new()); - - if db.needs_login() { - let password = - rbw::pinentry::getpin("prompt", "desc", tty).await.unwrap(); - let (access_token, refresh_token, iterations, protected_key, keys) = - rbw::actions::login(&email, &password).await.unwrap(); - - state.priv_key = Some(keys); - - db.access_token = Some(access_token); - db.refresh_token = Some(refresh_token); - db.iterations = Some(iterations); - db.protected_key = Some(protected_key); - db.save_async(&email).await.unwrap(); - } - - send_response(sock, &rbw::agent::Response::Ack).await; -} - -async fn unlock( - sock: &mut tokio::net::UnixStream, - state: std::sync::Arc>, - tty: Option<&str>, -) { - let mut state = state.write().await; - - if state.needs_unlock() { - let email = config_email().await; - let password = - rbw::pinentry::getpin("prompt", "desc", tty).await.unwrap(); - - let db = rbw::db::Db::load_async(&email) - .await - .unwrap_or_else(|_| rbw::db::Db::new()); - - let keys = rbw::actions::unlock( - &email, - &password, - db.iterations.unwrap(), - db.protected_key.as_deref().unwrap(), - ) - .await - .unwrap(); - - state.priv_key = Some(keys); - } - - send_response(sock, &rbw::agent::Response::Ack).await; -} - -async fn lock( - sock: &mut tokio::net::UnixStream, - state: std::sync::Arc>, -) { - let mut state = state.write().await; - - state.priv_key = None; - - send_response(sock, &rbw::agent::Response::Ack).await; -} - -async fn sync(sock: &mut tokio::net::UnixStream) { - let email = config_email().await; - let mut db = rbw::db::Db::load_async(&email) - .await - .unwrap_or_else(|_| rbw::db::Db::new()); - - let (protected_key, ciphers) = - rbw::actions::sync(db.access_token.as_ref().unwrap()) - .await - .unwrap(); - db.protected_key = Some(protected_key); - db.ciphers = ciphers; - db.save_async(&email).await.unwrap(); - - send_response(sock, &rbw::agent::Response::Ack).await; -} - -async fn decrypt( - sock: &mut tokio::net::UnixStream, - state: std::sync::Arc>, - cipherstring: &str, -) { - let state = state.read().await; - let keys = state.priv_key.as_ref().unwrap(); - let cipherstring = - rbw::cipherstring::CipherString::new(cipherstring).unwrap(); - let plaintext = - String::from_utf8(cipherstring.decrypt(keys).unwrap()).unwrap(); - - send_response(sock, &rbw::agent::Response::Decrypt { plaintext }).await; -} - -async fn handle_sock( - sock: tokio::net::UnixStream, - state: std::sync::Arc>, -) { - let mut buf = tokio::io::BufStream::new(sock); - let mut line = String::new(); - buf.read_line(&mut line).await.unwrap(); - let mut sock = buf.into_inner(); - let msg: rbw::agent::Request = serde_json::from_str(&line).unwrap(); - match msg.action { - rbw::agent::Action::Login => { - login(&mut sock, state.clone(), msg.tty.as_deref()).await - } - rbw::agent::Action::Unlock => { - unlock(&mut sock, state.clone(), msg.tty.as_deref()).await - } - rbw::agent::Action::Lock => lock(&mut sock, state.clone()).await, - rbw::agent::Action::Sync => sync(&mut sock).await, - rbw::agent::Action::Decrypt { cipherstring } => { - decrypt(&mut sock, state.clone(), &cipherstring).await - } - rbw::agent::Action::Quit => std::process::exit(0), - } -} - -async fn config_email() -> String { - let config = rbw::config::Config::load_async().await.unwrap(); - config.email.unwrap() -} - -struct State { - priv_key: Option, -} - -impl State { - fn needs_unlock(&self) -> bool { - self.priv_key.is_none() - } -} - -struct Agent { - timeout: tokio::time::Delay, - state: std::sync::Arc>, -} - -impl Agent { - fn new() -> Self { - let config = rbw::config::Config::load().unwrap(); - Self { - timeout: tokio::time::delay_for( - tokio::time::Duration::from_secs(config.lock_timeout), - ), - state: std::sync::Arc::new(tokio::sync::RwLock::new(State { - priv_key: None, - })), - } - } - - async fn run(&mut self, mut listener: tokio::net::UnixListener) { - loop { - tokio::select! { - sock = listener.next() => { - let state = self.state.clone(); - tokio::spawn(async move { - handle_sock(sock.unwrap().unwrap(), state).await - }); - } - _ = &mut self.timeout => { - self.state.write().await.priv_key = None - } - } - } - } -} - -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-agent/actions.rs b/src/bin/rbw-agent/actions.rs new file mode 100644 index 0000000..84e11bc --- /dev/null +++ b/src/bin/rbw-agent/actions.rs @@ -0,0 +1,116 @@ +pub async fn login( + sock: &mut crate::sock::Sock, + state: std::sync::Arc>, + tty: Option<&str>, +) { + let mut state = state.write().await; + let email = config_email().await; + let mut db = rbw::db::Db::load_async(&email) + .await + .unwrap_or_else(|_| rbw::db::Db::new()); + + if db.needs_login() { + let password = + rbw::pinentry::getpin("prompt", "desc", tty).await.unwrap(); + let (access_token, refresh_token, iterations, protected_key, keys) = + rbw::actions::login(&email, &password).await.unwrap(); + + state.priv_key = Some(keys); + + db.access_token = Some(access_token); + db.refresh_token = Some(refresh_token); + db.iterations = Some(iterations); + db.protected_key = Some(protected_key); + db.save_async(&email).await.unwrap(); + } + + respond_ack(sock).await; +} + +pub async fn unlock( + sock: &mut crate::sock::Sock, + state: std::sync::Arc>, + tty: Option<&str>, +) { + let mut state = state.write().await; + + if state.needs_unlock() { + let email = config_email().await; + let password = + rbw::pinentry::getpin("prompt", "desc", tty).await.unwrap(); + + let db = rbw::db::Db::load_async(&email) + .await + .unwrap_or_else(|_| rbw::db::Db::new()); + + let keys = rbw::actions::unlock( + &email, + &password, + db.iterations.unwrap(), + db.protected_key.as_deref().unwrap(), + ) + .await + .unwrap(); + + state.priv_key = Some(keys); + } + + respond_ack(sock).await; +} + +pub async fn lock( + sock: &mut crate::sock::Sock, + state: std::sync::Arc>, +) { + let mut state = state.write().await; + + state.priv_key = None; + + respond_ack(sock).await; +} + +pub async fn sync(sock: &mut crate::sock::Sock) { + let email = config_email().await; + let mut db = rbw::db::Db::load_async(&email) + .await + .unwrap_or_else(|_| rbw::db::Db::new()); + + let (protected_key, ciphers) = + rbw::actions::sync(db.access_token.as_ref().unwrap()) + .await + .unwrap(); + db.protected_key = Some(protected_key); + db.ciphers = ciphers; + db.save_async(&email).await.unwrap(); + + respond_ack(sock).await; +} + +pub async fn decrypt( + sock: &mut crate::sock::Sock, + state: std::sync::Arc>, + cipherstring: &str, +) { + let state = state.read().await; + let keys = state.priv_key.as_ref().unwrap(); + let cipherstring = + rbw::cipherstring::CipherString::new(cipherstring).unwrap(); + let plaintext = + String::from_utf8(cipherstring.decrypt(keys).unwrap()).unwrap(); + + respond_decrypt(sock, plaintext).await; +} + +async fn respond_ack(sock: &mut crate::sock::Sock) { + sock.send(&rbw::agent::Response::Ack).await; +} + +async fn respond_decrypt(sock: &mut crate::sock::Sock, plaintext: String) { + sock.send(&rbw::agent::Response::Decrypt { plaintext }) + .await; +} + +async fn config_email() -> String { + let config = rbw::config::Config::load_async().await.unwrap(); + config.email.unwrap() +} diff --git a/src/bin/rbw-agent/agent.rs b/src/bin/rbw-agent/agent.rs new file mode 100644 index 0000000..31e5ed8 --- /dev/null +++ b/src/bin/rbw-agent/agent.rs @@ -0,0 +1,74 @@ +use tokio::stream::StreamExt as _; + +pub struct State { + pub priv_key: Option, +} + +impl State { + pub fn needs_unlock(&self) -> bool { + self.priv_key.is_none() + } +} + +pub struct Agent { + timeout: tokio::time::Delay, + state: std::sync::Arc>, +} + +impl Agent { + pub fn new() -> Self { + let config = rbw::config::Config::load().unwrap(); + Self { + timeout: tokio::time::delay_for( + tokio::time::Duration::from_secs(config.lock_timeout), + ), + state: std::sync::Arc::new(tokio::sync::RwLock::new(State { + priv_key: None, + })), + } + } + + pub async fn run(&mut self, mut listener: tokio::net::UnixListener) { + loop { + tokio::select! { + sock = listener.next() => { + let mut sock + = crate::sock::Sock::new(sock.unwrap().unwrap()); + let state = self.state.clone(); + tokio::spawn(async move { + let req = sock.recv().await; + handle_request(&req, &mut sock, state.clone()).await; + }); + } + _ = &mut self.timeout => { + self.state.write().await.priv_key = None + } + } + } + } +} + +async fn handle_request( + req: &rbw::agent::Request, + sock: &mut crate::sock::Sock, + state: std::sync::Arc>, +) { + match &req.action { + rbw::agent::Action::Login => { + crate::actions::login(sock, state.clone(), req.tty.as_deref()) + .await + } + rbw::agent::Action::Unlock => { + crate::actions::unlock(sock, state.clone(), req.tty.as_deref()) + .await + } + rbw::agent::Action::Lock => { + crate::actions::lock(sock, state.clone()).await + } + rbw::agent::Action::Sync => crate::actions::sync(sock).await, + rbw::agent::Action::Decrypt { cipherstring } => { + crate::actions::decrypt(sock, state.clone(), &cipherstring).await + } + rbw::agent::Action::Quit => std::process::exit(0), + } +} diff --git a/src/bin/rbw-agent/daemon.rs b/src/bin/rbw-agent/daemon.rs new file mode 100644 index 0000000..64cf298 --- /dev/null +++ b/src/bin/rbw-agent/daemon.rs @@ -0,0 +1,50 @@ +pub struct StartupAck { + writer: std::os::unix::io::RawFd, +} + +impl StartupAck { + pub fn ack(&self) { + nix::unistd::write(self.writer, &[0]).unwrap(); + nix::unistd::close(self.writer).unwrap(); + } +} + +impl Drop for StartupAck { + fn drop(&mut self) { + nix::unistd::close(self.writer).unwrap(); + } +} + +pub fn daemonize() -> StartupAck { + 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), + } + } + } + + StartupAck { writer: w } +} diff --git a/src/bin/rbw-agent/main.rs b/src/bin/rbw-agent/main.rs new file mode 100644 index 0000000..235c4f2 --- /dev/null +++ b/src/bin/rbw-agent/main.rs @@ -0,0 +1,24 @@ +mod actions; +mod agent; +mod daemon; +mod sock; + +fn main() { + env_logger::from_env( + env_logger::Env::default().default_filter_or("info"), + ) + .init(); + + let startup_ack = daemon::daemonize(); + + // can't use tokio::main because we need to daemonize before starting the + // tokio runloop, or else things break + tokio::runtime::Runtime::new().unwrap().block_on(async { + let listener = crate::sock::listen(); + + startup_ack.ack(); + + let mut agent = crate::agent::Agent::new(); + agent.run(listener.unwrap()).await; + }) +} diff --git a/src/bin/rbw-agent/sock.rs b/src/bin/rbw-agent/sock.rs new file mode 100644 index 0000000..73ad5c2 --- /dev/null +++ b/src/bin/rbw-agent/sock.rs @@ -0,0 +1,36 @@ +use tokio::io::{AsyncBufReadExt as _, AsyncWriteExt as _}; + +pub struct Sock(tokio::net::UnixStream); + +impl Sock { + pub fn new(s: tokio::net::UnixStream) -> Self { + Self(s) + } + + pub async fn send(&mut self, res: &rbw::agent::Response) { + let Self(sock) = self; + sock.write_all(serde_json::to_string(res).unwrap().as_bytes()) + .await + .unwrap(); + sock.write_all(b"\n").await.unwrap(); + } + + pub async fn recv(&mut self) -> rbw::agent::Request { + let Self(sock) = self; + let mut buf = tokio::io::BufStream::new(sock); + let mut line = String::new(); + buf.read_line(&mut line).await.unwrap(); + serde_json::from_str(&line).unwrap() + } +} + +pub fn listen() -> anyhow::Result { + let runtime_dir = rbw::dirs::runtime_dir(); + std::fs::create_dir_all(&runtime_dir)?; + + let path = runtime_dir.join("socket"); + std::fs::remove_file(&path)?; + let sock = tokio::net::UnixListener::bind(&path)?; + log::debug!("listening on socket {}", path.to_string_lossy()); + Ok(sock) +} -- cgit v1.2.3-54-g00ecf