use crate::config::{Config, ConfigSnapshot};
use crate::error::{Error, Result};
use std::fs;
use std::io;
use std::path::{Path, PathBuf};
use chrono::Utc;
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
#[cfg(all(test, feature = "impure-test"))]
use crate::config::HookEntry;
const TRUST_DIR_NAME: &str = "kabu/trusted";
const TRUST_VERSION: u32 = 1;
#[cfg(all(test, feature = "impure-test"))]
thread_local! {
static TEST_TRUST_DIR: std::cell::RefCell<Option<PathBuf>> = const { std::cell::RefCell::new(None) };
}
#[derive(Debug, Serialize, Deserialize)]
pub(crate) struct TrustEntry {
pub version: u32,
pub main_worktree_path: PathBuf,
pub trusted_at: String,
pub config_snapshot: ConfigSnapshot,
}
fn trust_dir() -> Result<PathBuf> {
#[cfg(all(test, feature = "impure-test"))]
{
if let Some(path) = TEST_TRUST_DIR.with(|dir| dir.borrow().clone()) {
return Ok(path);
}
}
if let Some(path) = std::env::var_os("KABU_TRUST_DIR") {
return Ok(PathBuf::from(path));
}
let base = dirs::data_dir().ok_or(Error::TrustStorageNotFound)?;
Ok(base
.join(TRUST_DIR_NAME)
.join(format!("v{}", TRUST_VERSION)))
}
#[cfg(all(test, feature = "impure-test"))]
fn set_test_trust_dir(path: PathBuf) {
TEST_TRUST_DIR.with(|dir| {
*dir.borrow_mut() = Some(path);
});
}
pub(crate) fn compute_hash(main_worktree_path: &Path, config: &Config) -> Result<String> {
let mut hasher = Sha256::new();
let canonical_path =
main_worktree_path
.canonicalize()
.map_err(|e| Error::TrustVerificationFailed {
message: format!("Failed to canonicalize worktree path: {}", e),
})?;
let path_str = canonical_path.to_string_lossy();
hasher.update(path_str.as_bytes());
hasher.update(b"\n");
let snapshot = ConfigSnapshot::from_config(config);
let snapshot_json =
serde_json::to_string(&snapshot).map_err(|e| Error::TrustFileSerialization {
message: format!("Failed to serialize config snapshot: {}", e),
})?;
hasher.update(snapshot_json.as_bytes());
let result = hasher.finalize();
Ok(result.iter().map(|b| format!("{:02x}", b)).collect())
}
fn main_worktree_dir_name(path: &Path) -> String {
use sha2::{Digest, Sha256};
let path_str = path.to_string_lossy();
let mut hasher = Sha256::new();
hasher.update(path_str.as_bytes());
let hash = hasher.finalize();
let hex: String = hash.iter().map(|b| format!("{:02x}", b)).collect();
hex[..16].to_string()
}
pub(crate) fn is_trusted(main_worktree_path: &Path, config: &Config) -> Result<bool> {
if !config.hooks.has_hooks() {
return Ok(true); }
let canonical_path =
main_worktree_path
.canonicalize()
.map_err(|e| Error::TrustVerificationFailed {
message: format!("Failed to canonicalize worktree path: {}", e),
})?;
let hash = compute_hash(&canonical_path, config)?;
let dir_name = main_worktree_dir_name(&canonical_path);
let trust_file = trust_dir()?.join(&dir_name).join(format!("{}.yaml", hash));
let content = match fs::read_to_string(&trust_file) {
Ok(content) => content,
Err(e) if e.kind() == io::ErrorKind::NotFound => return Ok(false),
Err(e) => return Err(e.into()),
};
let entry: TrustEntry =
serde_yaml::from_str(&content).map_err(|e| Error::TrustFileCorrupted {
message: e.to_string(),
})?;
if entry.main_worktree_path != canonical_path {
return Ok(false);
}
Ok(true)
}
pub(crate) fn read_trust_entry(main_worktree_path: &Path) -> Result<Option<TrustEntry>> {
let canonical_path =
main_worktree_path
.canonicalize()
.map_err(|e| Error::TrustVerificationFailed {
message: format!("Failed to canonicalize worktree path: {}", e),
})?;
let dir_name = main_worktree_dir_name(&canonical_path);
let trust_dir = trust_dir()?;
let repo_trust_dir = trust_dir.join(&dir_name);
if !repo_trust_dir.exists() {
return Ok(None);
}
match fs::read_dir(&repo_trust_dir) {
Ok(entries) => {
for entry in entries.flatten() {
let path = entry.path();
if path.is_file() && path.extension().is_some_and(|ext| ext == "yaml") {
let content = fs::read_to_string(&path)?;
let trust_entry: TrustEntry = serde_yaml::from_str(&content).map_err(|e| {
Error::TrustFileSerialization {
message: format!("Failed to parse trust file: {}", e),
}
})?;
if trust_entry.main_worktree_path != canonical_path {
return Ok(None);
}
return Ok(Some(trust_entry));
}
}
Ok(None)
}
Err(e) if e.kind() == io::ErrorKind::NotFound => Ok(None),
Err(e) => Err(e.into()),
}
}
pub(crate) fn trust(main_worktree_path: &Path, config: &Config) -> Result<()> {
if !config.hooks.has_hooks() {
return Ok(());
}
let canonical_path =
main_worktree_path
.canonicalize()
.map_err(|e| Error::TrustVerificationFailed {
message: format!("Failed to canonicalize worktree path: {}", e),
})?;
let trust_base_path = trust_dir()?;
let dir_name = main_worktree_dir_name(&canonical_path);
let trust_repo_path = trust_base_path.join(&dir_name);
fs::create_dir_all(&trust_repo_path)?;
if trust_repo_path.exists() {
for entry in fs::read_dir(&trust_repo_path)?.flatten() {
let path = entry.path();
if path.extension().is_some_and(|ext| ext == "yaml") {
let _ = fs::remove_file(&path);
}
}
}
let hash = compute_hash(&canonical_path, config)?;
let trust_file = trust_repo_path.join(format!("{}.yaml", hash));
let snapshot = ConfigSnapshot::from_config(config);
let entry = TrustEntry {
version: TRUST_VERSION,
main_worktree_path: canonical_path,
trusted_at: Utc::now().to_rfc3339(),
config_snapshot: snapshot,
};
let content = serde_yaml::to_string(&entry).map_err(|e| Error::TrustFileSerialization {
message: e.to_string(),
})?;
fs::write(&trust_file, content)?;
Ok(())
}
pub(crate) fn untrust(main_worktree_path: &Path, config: &Config) -> Result<bool> {
if !config.hooks.has_hooks() {
return Ok(false);
}
let canonical_path =
main_worktree_path
.canonicalize()
.map_err(|e| Error::TrustVerificationFailed {
message: format!("Failed to canonicalize worktree path: {}", e),
})?;
let hash = compute_hash(&canonical_path, config)?;
let dir_name = main_worktree_dir_name(&canonical_path);
let trust_file = trust_dir()?.join(&dir_name).join(format!("{}.yaml", hash));
if trust_file.exists() {
fs::remove_file(&trust_file)?;
Ok(true)
} else {
Ok(false)
}
}
pub(crate) fn list_trusted() -> Result<Vec<TrustEntry>> {
let trust_path = trust_dir()?;
if !trust_path.exists() {
return Ok(Vec::new());
}
let mut entries = Vec::new();
for dir_entry in fs::read_dir(&trust_path)?.flatten() {
let dir_path = dir_entry.path();
if dir_path.is_dir() {
for file_entry in fs::read_dir(&dir_path)?.flatten() {
let file_path = file_entry.path();
if file_path.extension().is_some_and(|e| e == "yaml") {
match fs::read_to_string(&file_path) {
Ok(content) => match serde_yaml::from_str::<TrustEntry>(&content) {
Ok(trust_entry) => entries.push(trust_entry),
Err(e) => {
eprintln!(
"Failed to parse trust file: {}\n {}",
file_path.display(),
e
);
}
},
Err(e) => {
eprintln!(
"Failed to read trust file: {}\n {}",
file_path.display(),
e
);
}
}
}
}
}
}
Ok(entries)
}
#[cfg(all(test, feature = "impure-test"))]
#[allow(clippy::unwrap_used)]
mod tests {
use super::*;
use crate::config::{AutoCd, Config, Hooks, Mkdir, OnSetupFailure, Ui, Worktree};
use std::sync::OnceLock;
use tempfile::TempDir;
fn init_test_data_dir() {
static DATA_DIR: OnceLock<TempDir> = OnceLock::new();
let dir = DATA_DIR.get_or_init(|| TempDir::new().unwrap());
set_test_trust_dir(dir.path().to_path_buf());
}
fn create_test_config() -> Config {
Config {
on_conflict: None,
auto_cd: AutoCd::default(),
worktree: Worktree {
path_template: None,
branch_template: None,
},
on_setup_failure: OnSetupFailure::default(),
ui: Ui::default(),
hooks: Hooks::default(),
mkdir: Vec::new(),
link: Vec::new(),
copy: Vec::new(),
}
}
#[test]
fn test_compute_hash_empty_config() {
let temp_dir = TempDir::new().unwrap();
let config = create_test_config();
let hash = compute_hash(temp_dir.path(), &config).unwrap();
assert!(!hash.is_empty());
assert_eq!(hash.len(), 64); }
#[test]
fn test_compute_hash_different_hooks() {
let temp_dir = TempDir::new().unwrap();
let mut config1 = create_test_config();
config1.hooks = Hooks {
pre_add: vec![HookEntry {
command: "echo 'test1'".to_string(),
description: None,
}],
..Default::default()
};
let mut config2 = create_test_config();
config2.hooks = Hooks {
pre_add: vec![HookEntry {
command: "echo 'test2'".to_string(),
description: None,
}],
..Default::default()
};
let hash1 = compute_hash(temp_dir.path(), &config1).unwrap();
let hash2 = compute_hash(temp_dir.path(), &config2).unwrap();
assert_ne!(hash1, hash2);
}
#[test]
fn test_compute_hash_different_mkdir() {
let temp_dir = TempDir::new().unwrap();
use std::path::PathBuf;
let mut config1 = create_test_config();
config1.mkdir = vec![Mkdir {
path: PathBuf::from("dir1"),
description: None,
}];
let mut config2 = create_test_config();
config2.mkdir = vec![Mkdir {
path: PathBuf::from("dir2"),
description: None,
}];
let hash1 = compute_hash(temp_dir.path(), &config1).unwrap();
let hash2 = compute_hash(temp_dir.path(), &config2).unwrap();
assert_ne!(hash1, hash2);
}
#[test]
fn test_compute_hash_different_worktree() {
let temp_dir = TempDir::new().unwrap();
let mut config1 = create_test_config();
config1.worktree.path_template = Some("../wt1".to_string());
let mut config2 = create_test_config();
config2.worktree.path_template = Some("../wt2".to_string());
let hash1 = compute_hash(temp_dir.path(), &config1).unwrap();
let hash2 = compute_hash(temp_dir.path(), &config2).unwrap();
assert_ne!(hash1, hash2);
}
#[test]
fn test_is_trusted_empty_hooks() {
let temp_dir = TempDir::new().unwrap();
let config = create_test_config();
assert!(is_trusted(temp_dir.path(), &config).unwrap());
}
#[test]
fn test_trust_and_untrust() {
init_test_data_dir();
let temp_dir = TempDir::new().unwrap();
let mut config = create_test_config();
config.hooks = Hooks {
pre_add: vec![HookEntry {
command: "echo 'test'".to_string(),
description: None,
}],
..Default::default()
};
assert!(!is_trusted(temp_dir.path(), &config).unwrap());
trust(temp_dir.path(), &config).unwrap();
assert!(is_trusted(temp_dir.path(), &config).unwrap());
assert!(untrust(temp_dir.path(), &config).unwrap());
assert!(!is_trusted(temp_dir.path(), &config).unwrap());
assert!(!untrust(temp_dir.path(), &config).unwrap());
}
#[test]
fn test_list_trusted_no_error() {
init_test_data_dir();
let _entries = list_trusted().unwrap();
}
#[test]
fn test_list_trusted_with_trusted_repo() {
init_test_data_dir();
let temp_dir = TempDir::new().unwrap();
let mut config = create_test_config();
config.hooks = Hooks {
pre_add: vec![HookEntry {
command: "echo 'test'".to_string(),
description: None,
}],
post_add: vec![HookEntry {
command: "npm install".to_string(),
description: None,
}],
..Default::default()
};
trust(temp_dir.path(), &config).unwrap();
let entries = list_trusted().unwrap();
let canonical_path = temp_dir
.path()
.canonicalize()
.unwrap_or_else(|_| temp_dir.path().to_path_buf());
let found = entries
.iter()
.any(|e| e.main_worktree_path == canonical_path);
assert!(found, "Trusted repository should be in the list");
untrust(temp_dir.path(), &config).unwrap();
}
#[test]
fn test_list_trusted_multiple_repos() {
init_test_data_dir();
let temp_dir1 = TempDir::new().unwrap();
let temp_dir2 = TempDir::new().unwrap();
let mut config1 = create_test_config();
config1.hooks = Hooks {
pre_add: vec![HookEntry {
command: "echo 'repo1'".to_string(),
description: None,
}],
..Default::default()
};
let mut config2 = create_test_config();
config2.hooks = Hooks {
pre_add: vec![HookEntry {
command: "echo 'repo2'".to_string(),
description: None,
}],
..Default::default()
};
trust(temp_dir1.path(), &config1).unwrap();
trust(temp_dir2.path(), &config2).unwrap();
let entries = list_trusted().unwrap();
let canonical_path1 = temp_dir1
.path()
.canonicalize()
.unwrap_or_else(|_| temp_dir1.path().to_path_buf());
let canonical_path2 = temp_dir2
.path()
.canonicalize()
.unwrap_or_else(|_| temp_dir2.path().to_path_buf());
let found1 = entries
.iter()
.any(|e| e.main_worktree_path == canonical_path1);
let found2 = entries
.iter()
.any(|e| e.main_worktree_path == canonical_path2);
assert!(found1, "First trusted repository should be in the list");
assert!(found2, "Second trusted repository should be in the list");
untrust(temp_dir1.path(), &config1).unwrap();
untrust(temp_dir2.path(), &config2).unwrap();
}
#[test]
fn test_trust_entry_contains_main_worktree_path() {
init_test_data_dir();
let temp_dir = TempDir::new().unwrap();
let mut config = create_test_config();
config.hooks = Hooks {
hook_shell: None,
pre_add: vec![HookEntry {
command: "echo 'pre'".to_string(),
description: None,
}],
post_add: vec![HookEntry {
command: "npm install".to_string(),
description: None,
}],
pre_remove: vec![HookEntry {
command: "echo 'cleanup'".to_string(),
description: None,
}],
post_remove: vec![HookEntry {
command: "./scripts/cleanup.sh".to_string(),
description: None,
}],
};
trust(temp_dir.path(), &config).unwrap();
let entries = list_trusted().unwrap();
let canonical_path = temp_dir
.path()
.canonicalize()
.unwrap_or_else(|_| temp_dir.path().to_path_buf());
let entry = entries
.iter()
.find(|e| e.main_worktree_path == canonical_path)
.expect("Should find trusted entry");
assert_eq!(entry.main_worktree_path, canonical_path);
assert!(!entry.trusted_at.is_empty());
untrust(temp_dir.path(), &config).unwrap();
}
#[test]
fn test_is_trusted_hooks_changed() {
init_test_data_dir();
let temp_dir = TempDir::new().unwrap();
let mut config1 = create_test_config();
config1.hooks = Hooks {
pre_add: vec![HookEntry {
command: "echo 'original'".to_string(),
description: None,
}],
..Default::default()
};
trust(temp_dir.path(), &config1).unwrap();
assert!(is_trusted(temp_dir.path(), &config1).unwrap());
let mut config2 = create_test_config();
config2.hooks = Hooks {
pre_add: vec![HookEntry {
command: "echo 'modified'".to_string(),
description: None,
}],
..Default::default()
};
assert!(!is_trusted(temp_dir.path(), &config2).unwrap());
untrust(temp_dir.path(), &config1).unwrap();
}
#[test]
fn test_is_trusted_hooks_removed() {
init_test_data_dir();
let temp_dir = TempDir::new().unwrap();
let mut config = create_test_config();
config.hooks = Hooks {
pre_add: vec![HookEntry {
command: "echo 'test'".to_string(),
description: None,
}],
..Default::default()
};
trust(temp_dir.path(), &config).unwrap();
assert!(is_trusted(temp_dir.path(), &config).unwrap());
let empty_config = create_test_config();
assert!(is_trusted(temp_dir.path(), &empty_config).unwrap());
untrust(temp_dir.path(), &config).unwrap();
}
#[test]
fn test_is_trusted_different_main_worktree_path() {
use std::fs;
init_test_data_dir();
let temp_dir1 = TempDir::new().unwrap();
let temp_dir2 = TempDir::new().unwrap();
let mut config = create_test_config();
config.hooks = Hooks {
pre_add: vec![HookEntry {
command: "echo 'test'".to_string(),
description: None,
}],
..Default::default()
};
trust(temp_dir1.path(), &config).unwrap();
assert!(is_trusted(temp_dir1.path(), &config).unwrap());
let hash = compute_hash(temp_dir1.path(), &config).unwrap();
let dir_name = main_worktree_dir_name(
&temp_dir1
.path()
.canonicalize()
.unwrap_or_else(|_| temp_dir1.path().to_path_buf()),
);
let trust_subdir = trust_dir().unwrap().join(&dir_name);
fs::create_dir_all(&trust_subdir).unwrap();
let trust_file = trust_subdir.join(format!("{}.yaml", hash));
let snapshot = ConfigSnapshot::from_config(&config);
let fake_entry = TrustEntry {
version: TRUST_VERSION,
main_worktree_path: temp_dir2.path().canonicalize().unwrap(),
trusted_at: chrono::Utc::now().to_rfc3339(),
config_snapshot: snapshot,
};
let content = serde_yaml::to_string(&fake_entry).unwrap();
fs::write(&trust_file, content).unwrap();
assert!(!is_trusted(temp_dir1.path(), &config).unwrap());
untrust(temp_dir1.path(), &config).unwrap();
}
#[test]
fn test_is_trusted_corrupted_trust_file() {
use std::fs;
init_test_data_dir();
let temp_dir = TempDir::new().unwrap();
let mut config = create_test_config();
config.hooks = Hooks {
pre_add: vec![HookEntry {
command: "echo 'test'".to_string(),
description: None,
}],
..Default::default()
};
trust(temp_dir.path(), &config).unwrap();
assert!(is_trusted(temp_dir.path(), &config).unwrap());
let hash = compute_hash(temp_dir.path(), &config).unwrap();
let canonical_path = temp_dir
.path()
.canonicalize()
.unwrap_or_else(|_| temp_dir.path().to_path_buf());
let dir_name = main_worktree_dir_name(&canonical_path);
let trust_file = trust_dir()
.unwrap()
.join(&dir_name)
.join(format!("{}.yaml", hash));
fs::write(&trust_file, "invalid yaml content {{{").unwrap();
let result = is_trusted(temp_dir.path(), &config);
assert!(result.is_err());
fs::remove_file(&trust_file).ok();
}
}