From a6f1fe2c10ea9f7bb08d2344df62fee4a996cd69 Mon Sep 17 00:00:00 2001 From: Ellie Huxtable Date: Thu, 1 Feb 2024 15:00:46 +0000 Subject: [PATCH] feat: reencrypt/rekey local store (#1662) * feat: add record re-encrypting * automatically re-encrypt store when logging in with a different key * fix * actually save the new key lmao * add rekey * save new key * decode bip key * "add test for sqlite store re encrypt" --- Cargo.toml | 1 - atuin-client/src/encryption.rs | 10 +- atuin-client/src/record/sqlite_store.rs | 120 +++++++++++++++++++++- atuin-client/src/record/store.rs | 2 + atuin/src/command/client.rs | 2 +- atuin/src/command/client/account.rs | 5 +- atuin/src/command/client/account/login.rs | 67 +++++++++--- atuin/src/command/client/store.rs | 3 + atuin/src/command/client/store/rekey.rs | 64 ++++++++++++ atuin/src/command/client/sync.rs | 2 +- 10 files changed, 251 insertions(+), 25 deletions(-) create mode 100644 atuin/src/command/client/store/rekey.rs diff --git a/Cargo.toml b/Cargo.toml index b87d5820..9c5b8bf5 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -11,7 +11,6 @@ members = [ resolver = "2" [workspace.package] -name = "atuin" version = "17.2.1" authors = ["Ellie Huxtable "] rust-version = "1.67" diff --git a/atuin-client/src/encryption.rs b/atuin-client/src/encryption.rs index f4031059..50aacc24 100644 --- a/atuin-client/src/encryption.rs +++ b/atuin-client/src/encryption.rs @@ -30,6 +30,13 @@ pub struct EncryptedHistory { pub nonce: Nonce, } +pub fn generate_encoded_key() -> Result<(Key, String)> { + let key = XSalsa20Poly1305::generate_key(&mut OsRng); + let encoded = encode_key(&key)?; + + Ok((key, encoded)) +} + pub fn new_key(settings: &Settings) -> Result { let path = settings.key_path.as_str(); let path = PathBuf::from(path); @@ -38,8 +45,7 @@ pub fn new_key(settings: &Settings) -> Result { bail!("key already exists! cannot overwrite"); } - let key = XSalsa20Poly1305::generate_key(&mut OsRng); - let encoded = encode_key(&key)?; + let (key, encoded) = generate_encoded_key()?; let mut file = fs::File::create(path)?; file.write_all(encoded.as_bytes())?; diff --git a/atuin-client/src/record/sqlite_store.rs b/atuin-client/src/record/sqlite_store.rs index e9d7ff59..8bf200c3 100644 --- a/atuin-client/src/record/sqlite_store.rs +++ b/atuin-client/src/record/sqlite_store.rs @@ -19,6 +19,7 @@ use atuin_common::record::{ }; use uuid::Uuid; +use super::encryption::PASETO_V4; use super::store::Store; #[derive(Debug, Clone)] @@ -106,6 +107,15 @@ impl SqliteStore { }, } } + + async fn load_all(&self) -> Result>> { + let res = sqlx::query("select * from store ") + .map(Self::query_row) + .fetch_all(&self.pool) + .await?; + + Ok(res) + } } #[async_trait] @@ -251,13 +261,58 @@ impl Store for SqliteStore { Ok(res) } + + /// Reencrypt every single item in this store with a new key + /// Be careful - this may mess with sync. + async fn re_encrypt(&self, old_key: &[u8; 32], new_key: &[u8; 32]) -> Result<()> { + // Load all the records + // In memory like some of the other code here + // This will never be called in a hot loop, and only under the following circumstances + // 1. The user has logged into a new account, with a new key. They are unlikely to have a + // lot of data + // 2. The user has encountered some sort of issue, and runs a maintenance command that + // invokes this + let all = self.load_all().await?; + + let re_encrypted = all + .into_iter() + .map(|record| record.re_encrypt::(old_key, new_key)) + .collect::>>()?; + + // next up, we delete all the old data and reinsert the new stuff + // do it in one transaction, so if anything fails we rollback OK + + let mut tx = self.pool.begin().await?; + + let res = sqlx::query("delete from store").execute(&mut *tx).await?; + + let rows = res.rows_affected(); + debug!("deleted {rows} rows"); + + // don't call push_batch, as it will start its own transaction + // call the underlying save_raw + + for record in re_encrypted { + Self::save_raw(&mut tx, &record).await?; + } + + tx.commit().await?; + + Ok(()) + } } #[cfg(test)] mod tests { - use atuin_common::record::{EncryptedData, Host, HostId, Record}; + use atuin_common::{ + record::{DecryptedData, EncryptedData, Host, HostId, Record}, + utils::uuid_v7, + }; - use crate::record::{encryption::PASETO_V4, store::Store}; + use crate::{ + encryption::generate_encoded_key, + record::{encryption::PASETO_V4, store::Store}, + }; use super::SqliteStore; @@ -435,4 +490,65 @@ mod tests { "failed to insert 10k records" ); } + + #[tokio::test] + async fn re_encrypt() { + let store = SqliteStore::new(":memory:", 0.1).await.unwrap(); + let (key, _) = generate_encoded_key().unwrap(); + let data = vec![0u8, 1u8, 2u8, 3u8]; + let host_id = HostId(uuid_v7()); + + for i in 0..10 { + let record = Record::builder() + .host(Host::new(host_id)) + .version(String::from("test")) + .tag(String::from("test")) + .idx(i) + .data(DecryptedData(data.clone())) + .build(); + + let record = record.encrypt::(&key.into()); + store + .push(&record) + .await + .expect("failed to push encrypted record"); + } + + // first, check that we can decrypt the data with the current key + let all = store.all_tagged("test").await.unwrap(); + + assert_eq!(all.len(), 10, "failed to fetch all records"); + + for record in all { + let decrypted = record.decrypt::(&key.into()).unwrap(); + assert_eq!(decrypted.data.0, data); + } + + // reencrypt the store, then check if + // 1) it cannot be decrypted with the old key + // 2) it can be decrypted wiht the new key + + let (new_key, _) = generate_encoded_key().unwrap(); + store + .re_encrypt(&key.into(), &new_key.into()) + .await + .expect("failed to re-encrypt store"); + + let all = store.all_tagged("test").await.unwrap(); + + for record in all.iter() { + let decrypted = record.clone().decrypt::(&key.into()); + assert!( + decrypted.is_err(), + "did not get error decrypting with old key after re-encrypt" + ) + } + + for record in all { + let decrypted = record.decrypt::(&new_key.into()).unwrap(); + assert_eq!(decrypted.data.0, data); + } + + assert_eq!(store.len(host_id, "test").await.unwrap(), 10); + } } diff --git a/atuin-client/src/record/store.rs b/atuin-client/src/record/store.rs index 40c1224b..9c052213 100644 --- a/atuin-client/src/record/store.rs +++ b/atuin-client/src/record/store.rs @@ -28,6 +28,8 @@ pub trait Store { async fn last(&self, host: HostId, tag: &str) -> Result>>; async fn first(&self, host: HostId, tag: &str) -> Result>>; + async fn re_encrypt(&self, old_key: &[u8; 32], new_key: &[u8; 32]) -> Result<()>; + /// Get the next `limit` records, after and including the given index async fn next( &self, diff --git a/atuin/src/command/client.rs b/atuin/src/command/client.rs index 3ba5eb22..5b87d3ba 100644 --- a/atuin/src/command/client.rs +++ b/atuin/src/command/client.rs @@ -95,7 +95,7 @@ impl Cmd { Self::Sync(sync) => sync.run(settings, &db, sqlite_store).await, #[cfg(feature = "sync")] - Self::Account(account) => account.run(settings).await, + Self::Account(account) => account.run(settings, sqlite_store).await, Self::Kv(kv) => kv.run(&settings, &sqlite_store).await, diff --git a/atuin/src/command/client/account.rs b/atuin/src/command/client/account.rs index 75f8ed59..ee8e02a5 100644 --- a/atuin/src/command/client/account.rs +++ b/atuin/src/command/client/account.rs @@ -1,6 +1,7 @@ use clap::{Args, Subcommand}; use eyre::Result; +use atuin_client::record::sqlite_store::SqliteStore; use atuin_client::settings::Settings; pub mod change_password; @@ -33,9 +34,9 @@ pub enum Commands { } impl Cmd { - pub async fn run(self, settings: Settings) -> Result<()> { + pub async fn run(self, settings: Settings, store: SqliteStore) -> Result<()> { match self.command { - Commands::Login(l) => l.run(&settings).await, + Commands::Login(l) => l.run(&settings, &store).await, Commands::Register(r) => r.run(&settings).await, Commands::Logout => logout::run(&settings), Commands::Delete => delete::run(&settings).await, diff --git a/atuin/src/command/client/account/login.rs b/atuin/src/command/client/account/login.rs index 24f54ec2..9cd53399 100644 --- a/atuin/src/command/client/account/login.rs +++ b/atuin/src/command/client/account/login.rs @@ -6,7 +6,9 @@ use tokio::{fs::File, io::AsyncWriteExt}; use atuin_client::{ api_client, - encryption::{decode_key, encode_key, new_key, Key}, + encryption::{decode_key, encode_key, load_key, new_key, Key}, + record::sqlite_store::SqliteStore, + record::store::Store, settings::Settings, }; use atuin_common::api::LoginRequest; @@ -32,7 +34,7 @@ fn get_input() -> Result { } impl Cmd { - pub async fn run(&self, settings: &Settings) -> Result<()> { + pub async fn run(&self, settings: &Settings, store: &SqliteStore) -> Result<()> { let session_path = settings.session_path.as_str(); if PathBuf::from(session_path).exists() { @@ -44,24 +46,20 @@ impl Cmd { } let username = or_user_input(&self.username, "username"); - let key = or_user_input(&self.key, "encryption key [blank to use existing key file]"); let password = self.password.clone().unwrap_or_else(read_user_password); let key_path = settings.key_path.as_str(); - if key.is_empty() { - if PathBuf::from(key_path).exists() { - let bytes = fs_err::read_to_string(key_path) - .context("existing key file couldn't be read")?; - if decode_key(bytes).is_err() { - bail!("the key in existing key file was invalid"); - } - } else { - println!("No key file exists, creating a new"); - let _key = new_key(settings)?; - } + let key_path = PathBuf::from(key_path); + + let key = or_user_input(&self.key, "encryption key [blank to use existing key file]"); + + // if provided, the key may be EITHER base64, or a bip mnemonic + // try to normalize on base64 + let key = if key.is_empty() { + key } else { // try parse the key as a mnemonic... - let key = match bip39::Mnemonic::from_phrase(&key, bip39::Language::English) { + match bip39::Mnemonic::from_phrase(&key, bip39::Language::English) { Ok(mnemonic) => encode_key(Key::from_slice(mnemonic.entropy()))?, Err(err) => { if let Some(err) = err.downcast_ref::() { @@ -82,14 +80,51 @@ impl Cmd { key } } - }; + } + }; + // I've simplified this a little, but it could really do with a refactor + // Annoyingly, it's also very important to get it correct + if key.is_empty() { + if key_path.exists() { + let bytes = fs_err::read_to_string(key_path) + .context("existing key file couldn't be read")?; + if decode_key(bytes).is_err() { + bail!("the key in existing key file was invalid"); + } + } else { + println!("No key file exists, creating a new"); + let _key = new_key(settings)?; + } + } else if !key_path.exists() { if decode_key(key.clone()).is_err() { bail!("the specified key was invalid"); } let mut file = File::create(key_path).await?; file.write_all(key.as_bytes()).await?; + } else { + // we now know that the user has logged in specifying a key, AND that the key path + // exists + + // 1. check if the saved key and the provided key match. if so, nothing to do. + // 2. if not, re-encrypt the local history and overwrite the key + let current_key: [u8; 32] = load_key(settings)?.into(); + + let encoded = key.clone(); // gonna want to save it in a bit + let new_key: [u8; 32] = decode_key(key) + .context("could not decode provided key - is not valid base64")? + .into(); + + if new_key != current_key { + println!("\nRe-encrypting local store with new key"); + + store.re_encrypt(¤t_key, &new_key).await?; + + println!("Writing new key"); + let mut file = File::create(key_path).await?; + file.write_all(encoded.as_bytes()).await?; + } } let session = api_client::login( diff --git a/atuin/src/command/client/store.rs b/atuin/src/command/client/store.rs index 0d132c50..016f01b7 100644 --- a/atuin/src/command/client/store.rs +++ b/atuin/src/command/client/store.rs @@ -12,12 +12,14 @@ use time::OffsetDateTime; mod push; mod rebuild; +mod rekey; #[derive(Subcommand, Debug)] #[command(infer_subcommands = true)] pub enum Cmd { Status, Rebuild(rebuild::Rebuild), + Rekey(rekey::Rekey), #[cfg(feature = "sync")] Push(push::Push), @@ -33,6 +35,7 @@ impl Cmd { match self { Self::Status => self.status(store).await, Self::Rebuild(rebuild) => rebuild.run(settings, store, database).await, + Self::Rekey(rekey) => rekey.run(settings, store).await, #[cfg(feature = "sync")] Self::Push(push) => push.run(settings, store).await, diff --git a/atuin/src/command/client/store/rekey.rs b/atuin/src/command/client/store/rekey.rs new file mode 100644 index 00000000..3e079a5a --- /dev/null +++ b/atuin/src/command/client/store/rekey.rs @@ -0,0 +1,64 @@ +use clap::Args; +use eyre::{bail, Result}; +use tokio::{fs::File, io::AsyncWriteExt}; + +use atuin_client::{ + encryption::{decode_key, encode_key, generate_encoded_key, load_key, Key}, + record::sqlite_store::SqliteStore, + record::store::Store, + settings::Settings, +}; + +#[derive(Args, Debug)] +pub struct Rekey { + /// The new key to use for encryption. Omit for a randomly-generated key + key: Option, +} + +impl Rekey { + pub async fn run(&self, settings: &Settings, store: SqliteStore) -> Result<()> { + let key = if let Some(key) = self.key.clone() { + println!("Re-encrypting store with specified key"); + + let key = match bip39::Mnemonic::from_phrase(&key, bip39::Language::English) { + Ok(mnemonic) => encode_key(Key::from_slice(mnemonic.entropy()))?, + Err(err) => { + if let Some(err) = err.downcast_ref::() { + match err { + // assume they copied in the base64 key + bip39::ErrorKind::InvalidWord => key, + bip39::ErrorKind::InvalidChecksum => { + bail!("key mnemonic was not valid") + } + bip39::ErrorKind::InvalidKeysize(_) + | bip39::ErrorKind::InvalidWordLength(_) + | bip39::ErrorKind::InvalidEntropyLength(_, _) => { + bail!("key was not the correct length") + } + } + } else { + // unknown error. assume they copied the base64 key + key + } + } + }; + + key + } else { + println!("Re-encrypting store with freshly-generated key"); + let (_, encoded) = generate_encoded_key()?; + encoded + }; + + let current_key: [u8; 32] = load_key(settings)?.into(); + let new_key: [u8; 32] = decode_key(key.clone())?.into(); + + store.re_encrypt(¤t_key, &new_key).await?; + + println!("Store rewritten. Saving new key"); + let mut file = File::create(settings.key_path.clone()).await?; + file.write_all(key.as_bytes()).await?; + + Ok(()) + } +} diff --git a/atuin/src/command/client/sync.rs b/atuin/src/command/client/sync.rs index 5b438453..a7d57158 100644 --- a/atuin/src/command/client/sync.rs +++ b/atuin/src/command/client/sync.rs @@ -51,7 +51,7 @@ impl Cmd { ) -> Result<()> { match self { Self::Sync { force } => run(&settings, force, db, store).await, - Self::Login(l) => l.run(&settings).await, + Self::Login(l) => l.run(&settings, &store).await, Self::Logout => account::logout::run(&settings), Self::Register(r) => r.run(&settings).await, Self::Status => status::run(&settings, db).await,