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

@@ -11,7 +11,6 @@ members = [
 resolver = "2"
 
 [workspace.package]
-name = "atuin"
 version = "17.2.1"
 authors = ["Ellie Huxtable <ellie@elliehuxtable.com>"]
 rust-version = "1.67"

atuin-client/src/encryption.rs

@@ -30,6 +30,13 @@ pub struct EncryptedHistory {
     pub nonce: Nonce<XSalsa20Poly1305>,
 }
 
+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<Key> {
     let path = settings.key_path.as_str();
     let path = PathBuf::from(path);
@@ -38,8 +45,7 @@ pub fn new_key(settings: &Settings) -> Result<Key> {
         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())?;

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<Vec<Record<EncryptedData>>> {
+        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::<PASETO_V4>(old_key, new_key))
+            .collect::<Result<Vec<_>>>()?;
+
+        // 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::<PASETO_V4>(&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::<PASETO_V4>(&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::<PASETO_V4>(&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::<PASETO_V4>(&new_key.into()).unwrap();
+            assert_eq!(decrypted.data.0, data);
+        }
+
+        assert_eq!(store.len(host_id, "test").await.unwrap(), 10);
+    }
 }

atuin-client/src/record/store.rs

@@ -28,6 +28,8 @@ pub trait Store {
     async fn last(&self, host: HostId, tag: &str) -> Result<Option<Record<EncryptedData>>>;
     async fn first(&self, host: HostId, tag: &str) -> Result<Option<Record<EncryptedData>>>;
 
+    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,

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,
 

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,

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<String> {
 }
 
 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::<bip39::ErrorKind>() {
@@ -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(&current_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(

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,

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<String>,
+}
+
+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::<bip39::ErrorKind>() {
+                        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(&current_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(())
+    }
+}

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,