/// Store for shell vars /// I should abstract this and reuse code between the alias/env stores /// This is easier for now /// Once I have two implementations, building a common base is much easier. use std::collections::BTreeMap; use atuin_client::record::sqlite_store::SqliteStore; use atuin_common::record::{DecryptedData, Host, HostId}; use eyre::{Result, bail, ensure, eyre}; use atuin_client::record::encryption::PASETO_V4; use atuin_client::record::store::Store; use crate::shell::Var; const DOTFILES_VAR_VERSION: &str = "v0"; const DOTFILES_VAR_TAG: &str = "dotfiles-var"; const DOTFILES_VAR_LEN: usize = 20000; // 20kb max total len, way more than should be needed. #[derive(Debug, Clone, PartialEq, Eq)] pub enum VarRecord { Create(Var), // create a full record Delete(String), // delete by name } impl VarRecord { pub fn serialize(&self) -> Result { use rmp::encode; let mut output = vec![]; match self { VarRecord::Create(env) => { encode::write_u8(&mut output, 0)?; // create env.serialize(&mut output)?; } VarRecord::Delete(env) => { encode::write_u8(&mut output, 1)?; // delete encode::write_array_len(&mut output, 1)?; // 1 field encode::write_str(&mut output, env.as_str())?; } } Ok(DecryptedData(output)) } pub fn deserialize(data: &DecryptedData, version: &str) -> Result { use rmp::decode; fn error_report(err: E) -> eyre::Report { eyre!("{err:?}") } match version { DOTFILES_VAR_VERSION => { let mut bytes = decode::Bytes::new(&data.0); let record_type = decode::read_u8(&mut bytes).map_err(error_report)?; match record_type { // create 0 => { let env = Var::deserialize(&mut bytes)?; Ok(VarRecord::Create(env)) } // delete 1 => { let nfields = decode::read_array_len(&mut bytes).map_err(error_report)?; ensure!( nfields == 1, "too many entries in v0 dotfiles var delete record" ); let bytes = bytes.remaining_slice(); let (key, bytes) = decode::read_str_from_slice(bytes).map_err(error_report)?; if !bytes.is_empty() { bail!("trailing bytes in encoded dotfiles var record. malformed") } Ok(VarRecord::Delete(key.to_owned())) } n => { bail!("unknown Dotfiles var record type {n}") } } } _ => { bail!("unknown version {version:?}") } } } } #[derive(Debug, Clone)] pub struct VarStore { pub store: SqliteStore, pub host_id: HostId, pub encryption_key: [u8; 32], } impl VarStore { // will want to init the actual kv store when that is done pub fn new(store: SqliteStore, host_id: HostId, encryption_key: [u8; 32]) -> VarStore { VarStore { store, host_id, encryption_key, } } /// Escape a value for use in POSIX shells (bash, zsh) /// This adds double quotes around the value and escapes any embedded double quotes fn escape_posix_value(value: &str) -> String { // If the value contains no special characters, we can use it unquoted if value .chars() .all(|c| c.is_alphanumeric() || c == '_' || c == '-' || c == '/' || c == '.') { value.to_string() } else { // Otherwise, wrap in double quotes and escape any special characters format!( "\"{}\"", value .replace('\\', "\\\\") .replace('"', "\\\"") .replace('$', "\\$") .replace('`', "\\`") ) } } /// Escape a value for use in fish shell /// Fish uses single quotes for literal strings, but we need to handle embedded single quotes fn escape_fish_value(value: &str) -> String { // If the value contains no special characters, we can use it unquoted if value .chars() .all(|c| c.is_alphanumeric() || c == '_' || c == '-' || c == '/' || c == '.') { value.to_string() } else { // Use single quotes and escape any embedded single quotes format!("'{}'", value.replace('\'', "\\'")) } } /// Escape a value for use in xonsh /// Xonsh uses Python-style string literals fn escape_xonsh_value(value: &str) -> String { // If the value contains no special characters, we can use it unquoted if value .chars() .all(|c| c.is_alphanumeric() || c == '_' || c == '-' || c == '/' || c == '.') { value.to_string() } else { // Use double quotes and escape appropriately for Python strings format!("\"{}\"", value.replace('\\', "\\\\").replace('"', "\\\"")) } } pub async fn xonsh(&self) -> Result { let env = self.vars().await?; Ok(Self::format_xonsh(&env)) } pub async fn fish(&self) -> Result { let env = self.vars().await?; Ok(Self::format_fish(&env)) } pub async fn posix(&self) -> Result { let env = self.vars().await?; Ok(Self::format_posix(&env)) } pub async fn powershell(&self) -> Result { let env = self.vars().await?; Ok(Self::format_powershell(&env)) } fn format_xonsh(env: &[Var]) -> String { let mut config = String::new(); for env in env { let escaped_value = Self::escape_xonsh_value(&env.value); config.push_str(&format!("${}={}\n", env.name, escaped_value)); } config } fn format_fish(env: &[Var]) -> String { let mut config = String::new(); for env in env { let escaped_value = Self::escape_fish_value(&env.value); config.push_str(&format!("set -gx {} {}\n", env.name, escaped_value)); } config } fn format_posix(env: &[Var]) -> String { let mut config = String::new(); for env in env { let escaped_value = Self::escape_posix_value(&env.value); if env.export { config.push_str(&format!("export {}={}\n", env.name, escaped_value)); } else { config.push_str(&format!("{}={}\n", env.name, escaped_value)); } } config } fn format_powershell(env: &[Var]) -> String { let mut config = String::new(); for var in env { config.push_str(&crate::shell::powershell::format_var(var)); } config } pub async fn build(&self) -> Result<()> { let dir = atuin_common::utils::dotfiles_cache_dir(); tokio::fs::create_dir_all(dir.clone()).await?; let env = self.vars().await?; // Build for all supported shells let posix = Self::format_posix(&env); let xonsh = Self::format_xonsh(&env); let fsh = Self::format_fish(&env); let powershell = Self::format_powershell(&env); // All the same contents, maybe optimize in the future or perhaps there will be quirks // per-shell // I'd prefer separation atm let zsh = dir.join("vars.zsh"); let bash = dir.join("vars.bash"); let fish = dir.join("vars.fish"); let xsh = dir.join("vars.xsh"); let ps1 = dir.join("vars.ps1"); tokio::fs::write(zsh, &posix).await?; tokio::fs::write(bash, &posix).await?; tokio::fs::write(fish, &fsh).await?; tokio::fs::write(xsh, &xonsh).await?; tokio::fs::write(ps1, &powershell).await?; Ok(()) } pub async fn set(&self, name: &str, value: &str, export: bool) -> Result<()> { if name.len() + value.len() > DOTFILES_VAR_LEN { return Err(eyre!( "var record too large: max len {} bytes", DOTFILES_VAR_LEN )); } let record = VarRecord::Create(Var { name: name.to_string(), value: value.to_string(), export, }); let bytes = record.serialize()?; let idx = self .store .last(self.host_id, DOTFILES_VAR_TAG) .await? .map_or(0, |entry| entry.idx + 1); let record = atuin_common::record::Record::builder() .host(Host::new(self.host_id)) .version(DOTFILES_VAR_VERSION.to_string()) .tag(DOTFILES_VAR_TAG.to_string()) .idx(idx) .data(bytes) .build(); self.store .push(&record.encrypt::(&self.encryption_key)) .await?; // set mutates shell config, so build again self.build().await?; Ok(()) } pub async fn delete(&self, name: &str) -> Result<()> { if name.len() > DOTFILES_VAR_LEN { return Err(eyre!( "var record too large: max len {} bytes", DOTFILES_VAR_LEN, )); } let record = VarRecord::Delete(name.to_string()); let bytes = record.serialize()?; let idx = self .store .last(self.host_id, DOTFILES_VAR_TAG) .await? .map_or(0, |entry| entry.idx + 1); let record = atuin_common::record::Record::builder() .host(Host::new(self.host_id)) .version(DOTFILES_VAR_VERSION.to_string()) .tag(DOTFILES_VAR_TAG.to_string()) .idx(idx) .data(bytes) .build(); self.store .push(&record.encrypt::(&self.encryption_key)) .await?; // delete mutates shell config, so build again self.build().await?; Ok(()) } pub async fn vars(&self) -> Result> { let mut build = BTreeMap::new(); // this is sorted, oldest to newest let tagged = self.store.all_tagged(DOTFILES_VAR_TAG).await?; for record in tagged { let version = record.version.clone(); let decrypted = match version.as_str() { DOTFILES_VAR_VERSION => record.decrypt::(&self.encryption_key)?, version => bail!("unknown version {version:?}"), }; let ar = VarRecord::deserialize(&decrypted.data, version.as_str())?; match ar { VarRecord::Create(a) => { build.insert(a.name.clone(), a); } VarRecord::Delete(d) => { build.remove(&d); } } } Ok(build.into_values().collect()) } } #[cfg(test)] mod tests { use rand::rngs::OsRng; use atuin_client::record::sqlite_store::SqliteStore; use crate::{shell::Var, store::test_local_timeout}; use super::{DOTFILES_VAR_VERSION, VarRecord, VarStore}; use crypto_secretbox::{KeyInit, XSalsa20Poly1305}; #[test] fn encode_decode() { let record = Var { name: "BEEP".to_owned(), value: "boop".to_owned(), export: false, }; let record = VarRecord::Create(record); let snapshot = [ 204, 0, 147, 164, 66, 69, 69, 80, 164, 98, 111, 111, 112, 194, ]; let encoded = record.serialize().unwrap(); let decoded = VarRecord::deserialize(&encoded, DOTFILES_VAR_VERSION).unwrap(); assert_eq!(encoded.0, &snapshot); assert_eq!(decoded, record); } #[test] fn test_escape_posix_value() { // Simple values should not be quoted assert_eq!(VarStore::escape_posix_value("simple"), "simple"); assert_eq!(VarStore::escape_posix_value("path/to/file"), "path/to/file"); assert_eq!( VarStore::escape_posix_value("value_with_underscores"), "value_with_underscores" ); // Values with spaces should be quoted assert_eq!( VarStore::escape_posix_value("hello world"), "\"hello world\"" ); assert_eq!(VarStore::escape_posix_value("bar baz"), "\"bar baz\""); // Values with special characters should be quoted and escaped assert_eq!( VarStore::escape_posix_value("say \"hello\""), "\"say \\\"hello\\\"\"" ); assert_eq!( VarStore::escape_posix_value("path\\with\\backslashes"), "\"path\\\\with\\\\backslashes\"" ); assert_eq!( VarStore::escape_posix_value("say $hello"), "\"say \\$hello\"" ); assert_eq!( VarStore::escape_posix_value("see `example.md`"), "\"see \\`example.md\\`\"" ); } #[test] fn test_escape_fish_value() { // Simple values should not be quoted assert_eq!(VarStore::escape_fish_value("simple"), "simple"); assert_eq!(VarStore::escape_fish_value("path/to/file"), "path/to/file"); // Values with spaces should be single-quoted assert_eq!(VarStore::escape_fish_value("hello world"), "'hello world'"); assert_eq!(VarStore::escape_fish_value("bar baz"), "'bar baz'"); // Values with single quotes should be escaped assert_eq!(VarStore::escape_fish_value("don't"), "'don\\'t'"); } #[test] fn test_escape_xonsh_value() { // Simple values should not be quoted assert_eq!(VarStore::escape_xonsh_value("simple"), "simple"); assert_eq!(VarStore::escape_xonsh_value("path/to/file"), "path/to/file"); // Values with spaces should be quoted assert_eq!( VarStore::escape_xonsh_value("hello world"), "\"hello world\"" ); assert_eq!(VarStore::escape_xonsh_value("bar baz"), "\"bar baz\""); // Values with special characters should be quoted and escaped assert_eq!( VarStore::escape_xonsh_value("say \"hello\""), "\"say \\\"hello\\\"\"" ); assert_eq!( VarStore::escape_xonsh_value("path\\with\\backslashes"), "\"path\\\\with\\\\backslashes\"" ); } #[tokio::test] async fn build_vars() { let store = SqliteStore::new(":memory:", test_local_timeout()) .await .unwrap(); let key: [u8; 32] = XSalsa20Poly1305::generate_key(&mut OsRng).into(); let host_id = atuin_common::record::HostId(atuin_common::utils::uuid_v7()); let env = VarStore::new(store, host_id, key); env.set("BEEP", "boop", false).await.unwrap(); env.set("HOMEBREW_NO_AUTO_UPDATE", "1", true).await.unwrap(); let mut env_vars = env.vars().await.unwrap(); env_vars.sort_by_key(|a| a.name.clone()); assert_eq!(env_vars.len(), 2); assert_eq!( env_vars[0], Var { name: String::from("BEEP"), value: String::from("boop"), export: false, } ); assert_eq!( env_vars[1], Var { name: String::from("HOMEBREW_NO_AUTO_UPDATE"), value: String::from("1"), export: true, } ); } #[tokio::test] async fn test_var_generation_with_spaces() { let store = SqliteStore::new(":memory:", test_local_timeout()) .await .unwrap(); let key: [u8; 32] = XSalsa20Poly1305::generate_key(&mut OsRng).into(); let host_id = atuin_common::record::HostId(atuin_common::utils::uuid_v7()); let env = VarStore::new(store, host_id, key); // Test the exact scenario from the bug report env.set("FOO", "bar baz", true).await.unwrap(); let posix_output = env.posix().await.unwrap(); let fish_output = env.fish().await.unwrap(); let xonsh_output = env.xonsh().await.unwrap(); // POSIX should quote the value assert_eq!(posix_output, "export FOO=\"bar baz\"\n"); // Fish should quote the value assert_eq!(fish_output, "set -gx FOO 'bar baz'\n"); // Xonsh should quote the value assert_eq!(xonsh_output, "$FOO=\"bar baz\"\n"); } }