From 0172406624b0dbfdb893d2487d4d27a3aaa461a4 Mon Sep 17 00:00:00 2001 From: Filip Tibell Date: Fri, 22 Mar 2024 14:18:40 +0100 Subject: [PATCH] Implement trust subcommand and committing home --- Cargo.lock | 14 ++++++++ Cargo.toml | 1 + lib/storage/home.rs | 61 ++++++++++++++++++++++++---------- lib/storage/load_and_save.rs | 23 +++++++++++++ lib/storage/mod.rs | 1 + lib/storage/trust.rs | 34 +++++++++++-------- src/cli/debug_system_info.rs | 8 ++--- src/cli/debug_trusted_tools.rs | 10 +++--- src/cli/list.rs | 4 ++- src/cli/mod.rs | 29 +++++++++++----- src/cli/trust.rs | 40 ++++++++++++++++++++++ src/main.rs | 12 ++++--- 12 files changed, 181 insertions(+), 56 deletions(-) create mode 100644 lib/storage/load_and_save.rs create mode 100644 src/cli/trust.rs diff --git a/Cargo.lock b/Cargo.lock index c18bea5..edf5517 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -35,6 +35,7 @@ dependencies = [ "anyhow", "clap", "command-group", + "dashmap", "dialoguer", "dirs", "futures", @@ -416,6 +417,19 @@ dependencies = [ "syn", ] +[[package]] +name = "dashmap" +version = "5.5.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "978747c1d849a7d2ee5e8adc0159961c48fb7e5db2f06af6723b80123bb53856" +dependencies = [ + "cfg-if", + "hashbrown 0.14.3", + "lock_api", + "once_cell", + "parking_lot_core", +] + [[package]] name = "deranged" version = "0.3.11" diff --git a/Cargo.toml b/Cargo.toml index ace0323..50002f0 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -12,6 +12,7 @@ default-run = "aftman" [dependencies] anyhow = "1.0" clap = { version = "4.5", features = ["derive"] } +dashmap = "5.5" dialoguer = "0.11" dirs = "5.0" futures = "0.3" diff --git a/lib/storage/home.rs b/lib/storage/home.rs index a41df24..23bd2db 100644 --- a/lib/storage/home.rs +++ b/lib/storage/home.rs @@ -1,9 +1,8 @@ use std::env::var; use std::path::{Path, PathBuf}; +use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::Arc; -use tokio::fs::read_to_string; - use super::{StorageError, StorageResult, TrustStorage}; /** @@ -17,49 +16,75 @@ use super::{StorageError, StorageResult, TrustStorage}; #[derive(Debug, Clone)] pub struct Home { path: Arc, + saved: Arc, + trust: TrustStorage, } impl Home { /** Creates a new `Home` from the given path. */ - fn from_path(path: impl Into) -> Self { - Self { - path: path.into().into(), - } + async fn load_from_path(path: impl Into) -> StorageResult { + let path: Arc = path.into().into(); + let saved = Arc::new(AtomicBool::new(false)); + + let trust = TrustStorage::load(&path).await?; + + Ok(Self { path, saved, trust }) } /** Creates a new `Home` from the environment. + This will read, and if necessary, create the Aftman home directory + and its contents - including trust storage, tools storage, etc. + If the `AFTMAN_ROOT` environment variable is set, this will use that as the home directory. Otherwise, it will use `$HOME/.aftman`. */ - pub fn from_env() -> StorageResult { + pub async fn load_from_env() -> StorageResult { Ok(match var("AFTMAN_ROOT") { - Ok(root_str) => Self::from_path(root_str), + Ok(root_str) => Self::load_from_path(root_str).await?, Err(_) => { let path = dirs::home_dir() .ok_or(StorageError::HomeNotFound)? .join(".aftman"); - Self::from_path(path) + Self::load_from_path(path).await? } }) } /** - Reads the trust storage for this `Home`. + Returns a reference to the `TrustStorage` for this `Home`. + */ + pub fn trust(&self) -> &TrustStorage { + &self.trust + } - This function will return an error if the trust storage file - cannot be read - if it does not exist, it will be created. + /** + Saves the contents of this `Home` to disk. */ - pub async fn trust_storage(&self) -> StorageResult { - let path = self.path.join("trusted.txt"); - match read_to_string(&path).await { - Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(TrustStorage::new()), - Err(e) => Err(e.into()), - Ok(s) => Ok(TrustStorage::from_str(s)), + pub async fn save(&self) -> StorageResult<()> { + self.trust.save(&self.path).await?; + self.saved.store(true, Ordering::SeqCst); + Ok(()) + } +} + +// Implement Drop with an error message if the Home was dropped +// without being saved - this should never happen since a Home +// should always be loaded once on startup and saved on shutdown +// in the CLI, but this detail may be missed during refactoring. +// In the future, if AsyncDrop ever becomes a thing, we can just +// force the save to happen in the Drop implementation instead. +impl Drop for Home { + fn drop(&mut self) { + if !self.saved.load(Ordering::SeqCst) { + tracing::error!( + "Aftman home was dropped without being saved!\ + \nChanges to trust, tools, and more may have been lost." + ) } } } diff --git a/lib/storage/load_and_save.rs b/lib/storage/load_and_save.rs new file mode 100644 index 0000000..fd95ec7 --- /dev/null +++ b/lib/storage/load_and_save.rs @@ -0,0 +1,23 @@ +use std::path::Path; + +use tokio::fs::{read_to_string, write}; + +use super::{StorageResult, TrustStorage}; + +const FILE_PATH_TRUST: &str = "trusted.txt"; + +impl TrustStorage { + pub(super) async fn load(home_path: impl AsRef) -> StorageResult { + let path = home_path.as_ref().join(FILE_PATH_TRUST); + match read_to_string(&path).await { + Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(TrustStorage::new()), + Err(e) => Err(e.into()), + Ok(s) => Ok(TrustStorage::from_str(s)), + } + } + + pub(super) async fn save(&self, home_path: impl AsRef) -> StorageResult<()> { + let path = home_path.as_ref().join(FILE_PATH_TRUST); + Ok(write(path, self.to_string()).await?) + } +} diff --git a/lib/storage/mod.rs b/lib/storage/mod.rs index 1cce08e..abde34e 100644 --- a/lib/storage/mod.rs +++ b/lib/storage/mod.rs @@ -1,4 +1,5 @@ mod home; +mod load_and_save; mod result; mod trust; diff --git a/lib/storage/trust.rs b/lib/storage/trust.rs index 6171638..ec2b929 100644 --- a/lib/storage/trust.rs +++ b/lib/storage/trust.rs @@ -1,16 +1,18 @@ #![allow(clippy::should_implement_trait)] #![allow(clippy::inherent_to_string)] -use std::{collections::BTreeSet, convert::Infallible, str::FromStr}; +use std::{convert::Infallible, str::FromStr, sync::Arc}; + +use dashmap::DashSet; use crate::tool::ToolId; /** Storage for trusted tool identifiers. */ -#[derive(Debug, Default)] +#[derive(Debug, Default, Clone)] pub struct TrustStorage { - tools: BTreeSet, + tools: Arc>, } impl TrustStorage { @@ -35,8 +37,10 @@ impl TrustStorage { .as_ref() .lines() .filter_map(|line| line.parse::().ok()) - .collect(); - Self { tools } + .collect::>(); + Self { + tools: Arc::new(tools), + } } /** @@ -44,7 +48,7 @@ impl TrustStorage { Returns `true` if the tool was added and not already trusted. */ - pub fn add_tool(&mut self, tool: ToolId) -> bool { + pub fn add_tool(&self, tool: ToolId) -> bool { self.tools.insert(tool) } @@ -53,8 +57,8 @@ impl TrustStorage { Returns `true` if the tool was previously trusted and has now been removed. */ - pub fn remove_tool(&mut self, tool: &ToolId) -> bool { - self.tools.remove(tool) + pub fn remove_tool(&self, tool: &ToolId) -> bool { + self.tools.remove(tool).is_some() } /** @@ -65,10 +69,12 @@ impl TrustStorage { } /** - Get an iterator over the tools in this `TrustStorage`. + Get a sorted copy of the trusted tools in this `TrustStorage`. */ - pub fn iter_tools(&self) -> impl Iterator { - self.tools.iter() + pub fn all_tools(&self) -> Vec { + let mut sorted_tools = self.tools.iter().map(|id| id.clone()).collect::>(); + sorted_tools.sort(); + sorted_tools } /** @@ -78,9 +84,9 @@ impl TrustStorage { */ pub fn to_string(&self) -> String { let mut contents = self - .tools - .iter() - .map(ToString::to_string) + .all_tools() + .into_iter() + .map(|id| id.to_string()) .collect::>() .join("\n"); contents.push('\n'); diff --git a/src/cli/debug_system_info.rs b/src/cli/debug_system_info.rs index 69f27ab..9f029f9 100644 --- a/src/cli/debug_system_info.rs +++ b/src/cli/debug_system_info.rs @@ -1,14 +1,14 @@ use anyhow::Result; use clap::Parser; -use aftman::system::Description; +use aftman::{storage::Home, system::Description}; /// Prints out information about the system detected by Aftman. #[derive(Debug, Parser)] -pub struct GetSystemInfoSubcommand {} +pub struct DebugSystemInfoSubcommand {} -impl GetSystemInfoSubcommand { - pub async fn run(&self) -> Result<()> { +impl DebugSystemInfoSubcommand { + pub async fn run(&self, _home: &Home) -> Result<()> { let desc = Description::current(); println!("Current system information:"); println!("{desc:#?}"); diff --git a/src/cli/debug_trusted_tools.rs b/src/cli/debug_trusted_tools.rs index 40b7e9b..e1dc3bd 100644 --- a/src/cli/debug_trusted_tools.rs +++ b/src/cli/debug_trusted_tools.rs @@ -5,14 +5,12 @@ use aftman::storage::Home; /// Prints out information about currently trusted tools. #[derive(Debug, Parser)] -pub struct GetTrustedToolsSubcommand {} +pub struct DebugTrustedToolsSubcommand {} -impl GetTrustedToolsSubcommand { - pub async fn run(&self) -> Result<()> { - let home = Home::from_env()?; - let storage = home.trust_storage().await?; +impl DebugTrustedToolsSubcommand { + pub async fn run(&self, home: &Home) -> Result<()> { println!("Trusted tools:"); - for tool in storage.iter_tools() { + for tool in home.trust().all_tools() { println!("{tool}"); } Ok(()) diff --git a/src/cli/list.rs b/src/cli/list.rs index 95ad519..76a2564 100644 --- a/src/cli/list.rs +++ b/src/cli/list.rs @@ -1,12 +1,14 @@ use anyhow::Result; use clap::Parser; +use aftman::storage::Home; + /// Lists all existing tools managed by Aftman. #[derive(Debug, Parser)] pub struct ListSubcommand {} impl ListSubcommand { - pub async fn run(&self) -> Result<()> { + pub async fn run(&self, _home: &Home) -> Result<()> { Ok(()) } } diff --git a/src/cli/mod.rs b/src/cli/mod.rs index b667778..c3e2c6d 100644 --- a/src/cli/mod.rs +++ b/src/cli/mod.rs @@ -1,13 +1,16 @@ +use aftman::storage::Home; use anyhow::Result; use clap::Parser; mod debug_system_info; mod debug_trusted_tools; mod list; +mod trust; -use self::debug_system_info::GetSystemInfoSubcommand; -use self::debug_trusted_tools::GetTrustedToolsSubcommand; +use self::debug_system_info::DebugSystemInfoSubcommand; +use self::debug_trusted_tools::DebugTrustedToolsSubcommand; use self::list::ListSubcommand; +use self::trust::TrustSubcommand; #[derive(Debug, Parser)] #[clap(author, version, about)] @@ -26,21 +29,29 @@ impl Args { pub enum Subcommand { // Hidden subcommands (for debugging) #[clap(hide = true)] - DebugSystemInfo(GetSystemInfoSubcommand), + DebugSystemInfo(DebugSystemInfoSubcommand), #[clap(hide = true)] - DebugTrustedTools(GetTrustedToolsSubcommand), + DebugTrustedTools(DebugTrustedToolsSubcommand), // Public subcommands List(ListSubcommand), + Trust(TrustSubcommand), } impl Subcommand { pub async fn run(self) -> Result<()> { - match self { + let home = Home::load_from_env().await?; + + let result = match self { // Hidden subcommands - Self::DebugSystemInfo(cmd) => cmd.run().await, - Self::DebugTrustedTools(cmd) => cmd.run().await, + Self::DebugSystemInfo(cmd) => cmd.run(&home).await, + Self::DebugTrustedTools(cmd) => cmd.run(&home).await, // Public subcommands - Self::List(cmd) => cmd.run().await, - } + Self::List(cmd) => cmd.run(&home).await, + Self::Trust(cmd) => cmd.run(&home).await, + }; + + home.save().await?; + + result } } diff --git a/src/cli/trust.rs b/src/cli/trust.rs new file mode 100644 index 0000000..957d70e --- /dev/null +++ b/src/cli/trust.rs @@ -0,0 +1,40 @@ +use anyhow::{bail, Result}; +use clap::Parser; + +use aftman::{storage::Home, tool::ToolId}; + +/// Mark the given tool(s) as being trusted. +#[derive(Debug, Parser)] +pub struct TrustSubcommand { + /// The tool(s) to mark as trusted. + pub tools: Vec, +} + +impl TrustSubcommand { + pub async fn run(self, home: &Home) -> Result<()> { + if self.tools.is_empty() { + bail!("Please provide at least one tool to trust."); + } + + let trust_storage = home.trust(); + let (new_tools, existing_tools) = self + .tools + .into_iter() + .partition::, _>(|tool| trust_storage.add_tool(tool.clone())); + + if !new_tools.is_empty() { + println!("Trusted new tools:"); + for tool in new_tools { + println!(" - {tool}"); + } + } + if !existing_tools.is_empty() { + println!("Already trusted tools:"); + for tool in existing_tools { + println!(" - {tool}"); + } + } + + Ok(()) + } +} diff --git a/src/main.rs b/src/main.rs index 5351387..544aa64 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,13 +1,14 @@ -use anyhow::Result; +use std::process::exit; + use clap::Parser; -use tracing::level_filters::LevelFilter; +use tracing::{error, level_filters::LevelFilter}; use tracing_subscriber::EnvFilter; mod cli; use cli::Args; #[tokio::main] -async fn main() -> Result<()> { +async fn main() { let tracing_env_filter = EnvFilter::builder() .with_default_directive(LevelFilter::INFO.into()) .from_env_lossy() @@ -24,5 +25,8 @@ async fn main() -> Result<()> { .without_time() .init(); - Args::parse().run().await + if let Err(e) = Args::parse().run().await { + error!("{e}"); + exit(1); + } }