diff options
| author | Benedikt Peetz <benedikt.peetz@b-peetz.de> | 2026-07-20 12:57:36 +0200 |
|---|---|---|
| committer | Benedikt Peetz <benedikt.peetz@b-peetz.de> | 2026-07-20 12:57:36 +0200 |
| commit | 73be69cd99f5a4784fe69f5d78f423c72e837284 (patch) | |
| tree | cb47246d2c7edd8cd828abb25a033b732d2289ae /crates/daemon | |
| parent | chore: Move more stuff out of atuin-client (diff) | |
| download | atuin-73be69cd99f5a4784fe69f5d78f423c72e837284.zip | |
chore: Commit
Diffstat (limited to 'crates/daemon')
35 files changed, 7171 insertions, 1220 deletions
diff --git a/crates/daemon/Cargo.toml b/crates/daemon/Cargo.toml index 6d5de681..07fd84d7 100644 --- a/crates/daemon/Cargo.toml +++ b/crates/daemon/Cargo.toml @@ -14,8 +14,6 @@ repository = { workspace = true } [dependencies] turtle-common = { workspace = true } async-trait = "0.1.58" -atuin-nucleo-matcher = { workspace = true } -atuin-nucleo = { workspace = true } axum = "0.8" base64 = "0.22" clap = { version = "4.5.7", features = ["derive"] } diff --git a/crates/daemon/build.rs b/crates/daemon/build.rs index 646b8588..5abf6b2a 100644 --- a/crates/daemon/build.rs +++ b/crates/daemon/build.rs @@ -6,7 +6,6 @@ use protox::prost::Message; fn main() -> Result<(), std::io::Error> { let proto_paths = [ "proto/history.proto", - "proto/search.proto", "proto/control.proto", "proto/semantic.proto", ]; diff --git a/crates/daemon/proto/search.proto b/crates/daemon/proto/search.proto deleted file mode 100644 index 6b84acbd..00000000 --- a/crates/daemon/proto/search.proto +++ /dev/null @@ -1,35 +0,0 @@ -syntax = "proto3"; -package search; - -enum FilterMode { - GLOBAL = 0; - HOST = 1; - SESSION = 2; - DIRECTORY = 3; - WORKSPACE = 4; - SESSION_PRELOAD = 5; -} - -message SearchContext { - string session_id = 1; - string cwd = 2; - string hostname = 3; - string host_id = 4; - optional string git_root = 5; -} - -message SearchRequest { - string query = 1; - uint64 query_id = 2; // Incrementing ID to match responses to queries - FilterMode filter_mode = 3; - SearchContext context = 4; -} - -message SearchResponse { - uint64 query_id = 1; // Echo back the query ID - repeated bytes ids = 2; -} - -service Search { - rpc Search(stream SearchRequest) returns (stream SearchResponse); -} diff --git a/crates/daemon/src/aclient/api_client.rs b/crates/daemon/src/aclient/api_client.rs new file mode 100644 index 00000000..666e6dce --- /dev/null +++ b/crates/daemon/src/aclient/api_client.rs @@ -0,0 +1,197 @@ +use std::env; +use std::time::Duration; + +use eyre::{Result, bail, eyre}; +use reqwest::{Response, StatusCode, Url, header::HeaderMap}; +use tracing::debug; +use uuid::Uuid; + +use turtle_common::{api::ErrorResponse, record::RecordStatus}; +use turtle_common::{ + api::{ATUIN_CARGO_VERSION, ATUIN_HEADER_VERSION, ATUIN_VERSION}, + record::{EncryptedData, HostId, Record, RecordIdx}, + tls::ensure_crypto_provider, +}; + +use semver::Version; + +static APP_USER_AGENT: &str = concat!("atuin/", env!("CARGO_PKG_VERSION"),); + +pub(crate) struct Client<'a> { + sync_addr: &'a str, + user_id: Uuid, + inner: reqwest::Client, +} + +fn make_url(address: &str, path: &str, user_id: Uuid) -> Result<String> { + let address = address.strip_suffix('/').unwrap_or(address); + + // `join()` expects a trailing `/` in order to join paths + // e.g. it treats `http://host:port/subdir` as a file called `subdir` + let address = &format!("{address}/api/v0/{user_id}/"); + + // passing a path with a leading `/` will cause `join()` to replace the entire URL path + let path = path.strip_prefix("/").unwrap_or(path); + + let url = Url::parse(address) + .map(|url| url.join(path))? + .map_err(|_| eyre!("invalid address"))?; + + Ok(url.to_string()) +} + +pub(crate) fn ensure_version(response: &Response) -> Result<bool> { + let version = response.headers().get(ATUIN_HEADER_VERSION); + + let version = if let Some(version) = version { + match version.to_str() { + Ok(v) => Version::parse(v), + Err(e) => bail!("failed to parse server version: {:?}", e), + } + } else { + bail!("Server not reporting its version: it is either too old or unhealthy"); + }?; + + // If the client is newer than the server + if version.major < ATUIN_VERSION.major { + println!( + "Atuin version mismatch! In order to successfully sync, the server needs to run a newer version of Atuin" + ); + println!("Client: {ATUIN_CARGO_VERSION}"); + println!("Server: {version}"); + + return Ok(false); + } + + Ok(true) +} + +async fn handle_resp_error(resp: Response) -> Result<Response> { + let status = resp.status(); + let url = resp.url().to_string(); + + if status == StatusCode::SERVICE_UNAVAILABLE { + bail!( + "Service unavailable: check https://status.atuin.sh (or get in touch with your host)" + ); + } + + if status == StatusCode::TOO_MANY_REQUESTS { + bail!("Rate limited; please wait before doing that again"); + } + + if !status.is_success() { + if let Ok(error) = resp.json::<ErrorResponse<'_>>().await { + let reason = error.reason; + + if status.is_client_error() { + bail!("Invalid request to the service at {url}, {status} - {reason}.") + } + + bail!( + "There was an error with the atuin sync service at {url}, server error {status}: {reason}.\nIf the problem persists, contact the host" + ) + } + + bail!( + "There was an error with the atuin sync service at {url}, Status {status:?}.\nIf the problem persists, contact the host" + ) + } + + Ok(resp) +} + +impl<'a> Client<'a> { + pub(crate) fn new( + sync_addr: &'a str, + connect_timeout: u64, + timeout: u64, + user_id: Uuid, + ) -> Result<Self> { + ensure_crypto_provider(); + let mut headers = HeaderMap::new(); + + // used for semver server check + headers.insert(ATUIN_HEADER_VERSION, ATUIN_CARGO_VERSION.parse()?); + + Ok(Client { + user_id, + sync_addr, + inner: reqwest::Client::builder() + .user_agent(APP_USER_AGENT) + .default_headers(headers) + .connect_timeout(Duration::new(connect_timeout, 0)) + .timeout(Duration::new(timeout, 0)) + .build()?, + }) + } + + pub(crate) async fn delete_store(&self) -> Result<()> { + let url = make_url(self.sync_addr, "/store", self.user_id)?; + let url = Url::parse(url.as_str())?; + + let resp = self.inner.delete(url).send().await?; + + handle_resp_error(resp).await?; + + Ok(()) + } + + pub(crate) async fn post_records(&self, records: &[Record<EncryptedData>]) -> Result<()> { + let url = make_url(self.sync_addr, "/record", self.user_id)?; + let url = Url::parse(url.as_str())?; + + debug!("uploading {} records to {url}", records.len()); + + let resp = self.inner.post(url).json(records).send().await?; + handle_resp_error(resp).await?; + + Ok(()) + } + + pub(crate) async fn next_records( + &self, + host: HostId, + tag: String, + start: RecordIdx, + count: u64, + ) -> Result<Vec<Record<EncryptedData>>> { + debug!("fetching record/s from host {}/{}/{}", host.0, tag, start); + + let url = make_url( + self.sync_addr, + &format!( + "/record/next?host={}&tag={}&count={}&start={}", + host.0, tag, count, start + ), + self.user_id, + )?; + + let url = Url::parse(url.as_str())?; + + let resp = self.inner.get(url).send().await?; + let resp = handle_resp_error(resp).await?; + + let records = resp.json::<Vec<Record<EncryptedData>>>().await?; + + Ok(records) + } + + pub(crate) async fn record_status(&self) -> Result<RecordStatus> { + let url = make_url(self.sync_addr, "/record", self.user_id)?; + let url = Url::parse(url.as_str())?; + + let resp = self.inner.get(url).send().await?; + let resp = handle_resp_error(resp).await?; + + if !ensure_version(&resp)? { + bail!("could not sync records due to version mismatch"); + } + + let index = resp.json().await?; + + debug!("got remote index {index:?}"); + + Ok(index) + } +} diff --git a/crates/daemon/src/aclient/database/mod.rs b/crates/daemon/src/aclient/database/mod.rs new file mode 100644 index 00000000..36049f80 --- /dev/null +++ b/crates/daemon/src/aclient/database/mod.rs @@ -0,0 +1,1314 @@ +use std::{ + env, + path::{Path, PathBuf}, + str::FromStr, +}; + +use crate::aclient::utils::setup_db; +use fs_err::{self as fs}; +use itertools::Itertools; +use sql_builder::{SqlBuilder, SqlName, bind::Bind, esc, quote}; +use sqlx::{ + Result, Row, + sqlite::{SqliteConnectOptions, SqliteJournalMode, SqlitePool, SqliteRow, SqliteSynchronous}, +}; +use time::OffsetDateTime; +use tracing::debug; +use turtle_common::utils; +use uuid::Uuid; + +use crate::aclient::{ + history::{HistoryId, HistoryStats}, + utils::get_host_user, +}; + +use super::{ + history::History, + ordering, + settings::{FilterMode, SearchMode, Settings}, +}; + +#[derive(Clone)] +pub(crate) struct Context { + pub(crate) session: String, + pub(crate) cwd: String, + pub(crate) hostname: String, + pub(crate) host_id: String, + pub(crate) git_root: Option<PathBuf>, +} + +#[derive(Default, Clone)] +pub(crate) struct OptFilters { + pub(crate) exit: Option<i64>, + pub(crate) exclude_exit: Option<i64>, + pub(crate) cwd: Option<String>, + pub(crate) exclude_cwd: Option<String>, + pub(crate) before: Option<String>, + pub(crate) after: Option<String>, + pub(crate) limit: Option<i64>, + pub(crate) offset: Option<i64>, + pub(crate) reverse: bool, + pub(crate) include_duplicates: bool, +} + +pub(crate) async fn current_context() -> eyre::Result<Context> { + let session = env::var("ATUIN_SESSION").map_err(|_| { + eyre::eyre!("Failed to find $ATUIN_SESSION in the environment. Check that you have correctly set up your shell.") + })?; + let hostname = get_host_user(); + let cwd = utils::get_current_dir(); + let host_id = Settings::host_id().await?; + let git_root = utils::in_git_repo(cwd.as_str()); + + Ok(Context { + session, + hostname, + cwd, + git_root, + host_id: host_id.0.as_simple().to_string(), + }) +} + +impl Context { + pub(crate) fn from_history(entry: &History) -> Self { + Self { + session: entry.session.clone(), + cwd: entry.cwd.clone(), + hostname: entry.hostname.clone(), + host_id: String::new(), + git_root: utils::in_git_repo(entry.cwd.as_str()), + } + } +} + +fn get_session_start_time(session_id: &str) -> Option<i64> { + if let Ok(uuid) = Uuid::parse_str(session_id) + && let Some(timestamp) = uuid.get_timestamp() + { + let (seconds, nanos) = timestamp.to_unix(); + return Some(seconds as i64 * 1_000_000_000 + i64::from(nanos)); + } + None +} + +// Intended for use on a developer machine and not a sync server. +// TODO: implement IntoIterator +#[derive(Debug, Clone)] +pub struct ClientSqlite { + pub(crate) pool: SqlitePool, +} + +impl ClientSqlite { + pub async fn new(path: impl AsRef<Path>, timeout: f64) -> Result<Self> { + fn mk_opts(path: &str) -> Result<SqliteConnectOptions> { + let opts = SqliteConnectOptions::from_str(path)? + .journal_mode(SqliteJournalMode::Wal) + .optimize_on_close(true, None) + .synchronous(SqliteSynchronous::Normal) + .with_regexp() + .create_if_missing(true); + + Ok(opts) + } + + let path = path.as_ref(); + debug!("opening sqlite database at {path:?}"); + + if utils::broken_symlink(path) { + eprintln!( + "Atuin: Sqlite db path ({}) is a broken symlink. Unable to read or create replacement.", + path.display() + ); + std::process::exit(1); + } + + if !path.exists() + && let Some(dir) = path.parent() + { + fs::create_dir_all(dir)?; + } + + let pool = setup_db!(path, timeout, mk_opts, "./db/client-migrations").await?; + Ok(Self { pool }) + } + + async fn save_raw(tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>, h: &History) -> Result<()> { + sqlx::query( + "insert or ignore into history(id, timestamp, duration, exit, command, cwd, session, hostname, author, intent, deleted_at) + values(?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11)", + ) + .bind(h.id.0.as_str()) + .bind(h.timestamp.unix_timestamp_nanos() as i64) + .bind(h.duration) + .bind(h.exit) + .bind(h.command.as_str()) + .bind(h.cwd.as_str()) + .bind(h.session.as_str()) + .bind(h.hostname.as_str()) + .bind(h.author.as_str()) + .bind(h.intent.as_deref()) + .bind(h.deleted_at.map(|t|t.unix_timestamp_nanos() as i64)) + .execute(&mut **tx) + .await?; + + Ok(()) + } + + async fn delete_row_raw( + tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>, + id: HistoryId, + ) -> Result<()> { + sqlx::query("delete from history where id = ?1") + .bind(id.0.as_str()) + .execute(&mut **tx) + .await?; + + Ok(()) + } + + #[expect(clippy::needless_pass_by_value)] + fn query_history_inner(row: SqliteRow) -> History { + let deleted_at: Option<i64> = row.get("deleted_at"); + let hostname: String = row.get("hostname"); + let author: Option<String> = row.try_get("author").ok().flatten(); + let author = author + .filter(|author| !author.trim().is_empty()) + .unwrap_or_else(|| History::author_from_hostname(hostname.as_str())); + let intent: Option<String> = row.try_get("intent").ok().flatten(); + let intent = intent.filter(|intent| !intent.trim().is_empty()); + + History::from_db() + .id(row.get("id")) + .timestamp( + OffsetDateTime::from_unix_timestamp_nanos(i128::from( + row.get::<i64, _>("timestamp"), + )) + .unwrap(), + ) + .duration(row.get("duration")) + .exit(row.get("exit")) + .command(row.get("command")) + .cwd(row.get("cwd")) + .session(row.get("session")) + .hostname(hostname) + .author(author) + .intent(intent) + .deleted_at( + deleted_at + .and_then(|t| OffsetDateTime::from_unix_timestamp_nanos(i128::from(t)).ok()), + ) + .build() + .into() + } +} + +impl ClientSqlite { + pub(crate) async fn save(&self, h: &History) -> Result<()> { + debug!("saving history to sqlite"); + let mut tx = self.pool.begin().await?; + Self::save_raw(&mut tx, h).await?; + tx.commit().await?; + + Ok(()) + } + + pub(crate) async fn save_bulk(&self, h: &[History]) -> Result<()> { + debug!("saving history to sqlite"); + + let mut tx = self.pool.begin().await?; + + for i in h { + Self::save_raw(&mut tx, i).await?; + } + + tx.commit().await?; + + Ok(()) + } + + pub(crate) async fn load(&self, id: &str) -> Result<Option<History>> { + debug!("loading history item {}", id); + + let res = sqlx::query("select * from history where id = ?1") + .bind(id) + .map(Self::query_history_inner) + .fetch_optional(&self.pool) + .await?; + + Ok(res) + } + + // make a unique list, that only shows the *newest* version of things + pub(crate) async fn list( + &self, + filters: &[FilterMode], + context: &Context, + max: Option<usize>, + unique: bool, + include_deleted: bool, + ) -> Result<Vec<History>> { + debug!("listing history"); + + let mut query = SqlBuilder::select_from(SqlName::new("history").alias("h").baquoted()); + query.field("*").order_desc("timestamp"); + if !include_deleted { + query.and_where_is_null("deleted_at"); + } + + let git_root = context.git_root.clone().map_or_else( + || context.cwd.clone(), + |git_root| git_root.to_str().unwrap_or("/").to_string(), + ); + + let session_start = get_session_start_time(&context.session); + + for filter in filters { + match filter { + FilterMode::Global => &mut query, + FilterMode::Host => query.and_where_eq("hostname", quote(&context.hostname)), + FilterMode::Session => query.and_where_eq("session", quote(&context.session)), + FilterMode::SessionPreload => { + query.and_where_eq("session", quote(&context.session)); + if let Some(session_start) = session_start { + query.or_where_lt("timestamp", session_start); + } + &mut query + } + FilterMode::Directory => query.and_where_eq("cwd", quote(&context.cwd)), + FilterMode::Workspace => query.and_where_like_left("cwd", &git_root), + }; + } + + if unique { + query.group_by("command").having("max(timestamp)"); + } + + if let Some(max) = max { + query.limit(max); + } + + let query = query.sql().expect("bug in list query. please report"); + + let res = sqlx::query(&query) + .map(Self::query_history_inner) + .fetch_all(&self.pool) + .await?; + + Ok(res) + } + + pub(crate) async fn range( + &self, + from: OffsetDateTime, + to: OffsetDateTime, + ) -> Result<Vec<History>> { + debug!("listing history from {:?} to {:?}", from, to); + + let res = sqlx::query( + "select * from history where timestamp >= ?1 and timestamp <= ?2 order by timestamp asc", + ) + .bind(from.unix_timestamp_nanos() as i64) + .bind(to.unix_timestamp_nanos() as i64) + .map(Self::query_history_inner) + .fetch_all(&self.pool) + .await?; + + Ok(res) + } + + pub(crate) async fn last(&self) -> Result<Option<History>> { + let res = sqlx::query( + "select * from history where duration >= 0 order by timestamp desc limit 1", + ) + .map(Self::query_history_inner) + .fetch_optional(&self.pool) + .await?; + + Ok(res) + } + + pub(crate) async fn history_count(&self, include_deleted: bool) -> Result<i64> { + let query = if include_deleted { + "select count(1) from history" + } else { + "select count(1) from history where deleted_at is null" + }; + + let res: (i64,) = sqlx::query_as(query).fetch_one(&self.pool).await?; + Ok(res.0) + } + + // Yes I know, it's a lot. + // Could maybe break it down to a searchparams struct or smth but that feels a little... pointless. + // Been debating maybe a DSL for search? eg "before:time limit:1 the query" + #[expect(clippy::too_many_lines)] + pub(crate) async fn search( + &self, + search_mode: SearchMode, + filter: FilterMode, + context: &Context, + query: &str, + filter_options: OptFilters, + ) -> Result<Vec<History>> { + let mut sql = SqlBuilder::select_from("history"); + + if !filter_options.include_duplicates { + sql.group_by("command").having("max(timestamp)"); + } + + if let Some(limit) = filter_options.limit { + sql.limit(limit); + } + + if let Some(offset) = filter_options.offset { + sql.offset(offset); + } + + if filter_options.reverse { + sql.order_asc("timestamp"); + } else { + sql.order_desc("timestamp"); + } + + let git_root = context.git_root.clone().map_or_else( + || context.cwd.clone(), + |git_root| git_root.to_str().unwrap_or("/").to_string(), + ); + + let session_start = get_session_start_time(&context.session); + + match filter { + FilterMode::Global => &mut sql, + FilterMode::Host => { + sql.and_where_eq("lower(hostname)", quote(context.hostname.to_lowercase())) + } + FilterMode::Session => sql.and_where_eq("session", quote(&context.session)), + FilterMode::SessionPreload => { + sql.and_where_eq("session", quote(&context.session)); + if let Some(session_start) = session_start { + sql.or_where_lt("timestamp", session_start); + } + &mut sql + } + FilterMode::Directory => sql.and_where_eq("cwd", quote(&context.cwd)), + FilterMode::Workspace => sql.and_where_like_left("cwd", git_root), + }; + + let orig_query = query; + + let mut regexes = Vec::new(); + if search_mode == SearchMode::Prefix { + sql.and_where_like_left("command", query.replace('*', "%")) + } else { + let mut is_or = false; + for token in QueryTokenizer::new(query) { + // TODO smart case mode could be made configurable like in fzf + let (is_glob, glob) = if token.has_uppercase() { + (true, "*") + } else { + (false, "%") + }; + let param = match token { + QueryToken::Regex(r) => { + regexes.push(String::from(r)); + continue; + } + QueryToken::Or => { + if !is_or { + is_or = true; + continue; + } + + format!("{glob}|{glob}") + } + QueryToken::MatchStart(term, _) => { + format!("{term}{glob}") + } + QueryToken::MatchEnd(term, _) => { + format!("{glob}{term}") + } + QueryToken::MatchFull(term, _) => { + format!("{glob}{term}{glob}") + } + QueryToken::Match(term, _) => { + if search_mode == SearchMode::FullText { + format!("{glob}{term}{glob}") + } else { + term.split("").join(glob) + } + } + }; + + sql.fuzzy_condition("command", param, token.is_inverse(), is_glob, is_or); + is_or = false; + } + + &mut sql + }; + + for regex in regexes { + sql.and_where("command regexp ?".bind(®ex)); + } + + filter_options + .exit + .map(|exit| sql.and_where_eq("exit", exit)); + + filter_options + .exclude_exit + .map(|exclude_exit| sql.and_where_ne("exit", exclude_exit)); + + filter_options + .cwd + .map(|cwd| sql.and_where_eq("cwd", quote(cwd))); + + filter_options + .exclude_cwd + .map(|exclude_cwd| sql.and_where_ne("cwd", quote(exclude_cwd))); + + filter_options.before.map(|before| { + interim::parse_date_string( + before.as_str(), + OffsetDateTime::now_utc(), + interim::Dialect::Uk, + ) + .map(|before| { + sql.and_where_lt("timestamp", quote(before.unix_timestamp_nanos() as i64)) + }) + }); + + filter_options.after.map(|after| { + interim::parse_date_string( + after.as_str(), + OffsetDateTime::now_utc(), + interim::Dialect::Uk, + ) + .map(|after| sql.and_where_gt("timestamp", quote(after.unix_timestamp_nanos() as i64))) + }); + + sql.and_where_is_null("deleted_at"); + + let query = sql.sql().expect("bug in search query. please report"); + + let res = sqlx::query(&query) + .map(Self::query_history_inner) + .fetch_all(&self.pool) + .await?; + + Ok(ordering::reorder_fuzzy(search_mode, orig_query, res)) + } + + pub(crate) async fn query_history(&self, query: &str) -> Result<Vec<History>> { + let res = sqlx::query(query) + .map(Self::query_history_inner) + .fetch_all(&self.pool) + .await?; + + Ok(res) + } + + pub(crate) async fn all_with_count(&self) -> Result<Vec<(History, i32)>> { + debug!("listing history"); + + let mut query = SqlBuilder::select_from(SqlName::new("history").alias("h").baquoted()); + + query + .fields(&[ + "id", + "max(timestamp) as timestamp", + "max(duration) as duration", + "exit", + "command", + "deleted_at", + "null as author", + "null as intent", + "group_concat(cwd, ':') as cwd", + "group_concat(session) as session", + "group_concat(hostname, ',') as hostname", + "count(*) as count", + ]) + .group_by("command") + .group_by("exit") + .and_where("deleted_at is null") + .order_desc("timestamp"); + + let query = query.sql().expect("bug in list query. please report"); + + let res = sqlx::query(&query) + .map(|row: SqliteRow| { + let count: i32 = row.get("count"); + (Self::query_history_inner(row), count) + }) + .fetch_all(&self.pool) + .await?; + + Ok(res) + } + + pub(crate) fn all_paged(&self, page_size: usize, include_deleted: bool, unique: bool) -> Paged { + Paged::new(self.clone(), page_size, include_deleted, unique) + } + + pub(crate) async fn delete_rows(&self, ids: &[HistoryId]) -> Result<()> { + let mut tx = self.pool.begin().await?; + + for id in ids { + Self::delete_row_raw(&mut tx, id.clone()).await?; + } + + tx.commit().await?; + + Ok(()) + } + + pub(crate) async fn stats(&self, h: &History) -> Result<HistoryStats> { + // We select the previous in the session by time + let mut prev = SqlBuilder::select_from("history"); + prev.field("*") + .and_where("timestamp < ?1") + .and_where("session = ?2") + .order_by("timestamp", true) + .limit(1); + + let mut next = SqlBuilder::select_from("history"); + next.field("*") + .and_where("timestamp > ?1") + .and_where("session = ?2") + .order_by("timestamp", false) + .limit(1); + + let mut total = SqlBuilder::select_from("history"); + total.field("count(1)").and_where("command = ?1"); + + let mut average = SqlBuilder::select_from("history"); + average.field("avg(duration)").and_where("command = ?1"); + + let mut exits = SqlBuilder::select_from("history"); + exits + .fields(&["exit", "count(1) as count"]) + .and_where("command = ?1") + .group_by("exit"); + + // rewrite the following with sqlbuilder + let mut day_of_week = SqlBuilder::select_from("history"); + day_of_week + .fields(&[ + "strftime('%w', ROUND(timestamp / 1000000000), 'unixepoch') AS day_of_week", + "count(1) as count", + ]) + .and_where("command = ?1") + .group_by("day_of_week"); + + // Intentionally format the string with 01 hardcoded. We want the average runtime for the + // _entire month_, but will later parse it as a datetime for sorting + // Sqlite has no datetime so we cannot do it there, and otherwise sorting will just be a + // string sort, which won't be correct. + let mut duration_over_time = SqlBuilder::select_from("history"); + duration_over_time + .fields(&[ + "strftime('01-%m-%Y', ROUND(timestamp / 1000000000), 'unixepoch') AS month_year", + "avg(duration) as duration", + ]) + .and_where("command = ?1") + .group_by("month_year") + .having("duration > 0"); + + let prev = prev.sql().expect("issue in stats previous query"); + let next = next.sql().expect("issue in stats next query"); + let total = total.sql().expect("issue in stats average query"); + let average = average.sql().expect("issue in stats previous query"); + let exits = exits.sql().expect("issue in stats exits query"); + let day_of_week = day_of_week.sql().expect("issue in stats day of week query"); + let duration_over_time = duration_over_time + .sql() + .expect("issue in stats duration over time query"); + + let prev = sqlx::query(&prev) + .bind(h.timestamp.unix_timestamp_nanos() as i64) + .bind(&h.session) + .map(Self::query_history_inner) + .fetch_optional(&self.pool) + .await?; + + let next = sqlx::query(&next) + .bind(h.timestamp.unix_timestamp_nanos() as i64) + .bind(&h.session) + .map(Self::query_history_inner) + .fetch_optional(&self.pool) + .await?; + + let total: (i64,) = sqlx::query_as(&total) + .bind(&h.command) + .fetch_one(&self.pool) + .await?; + + let average: (f64,) = sqlx::query_as(&average) + .bind(&h.command) + .fetch_one(&self.pool) + .await?; + + let exits: Vec<(i64, i64)> = sqlx::query_as(&exits) + .bind(&h.command) + .fetch_all(&self.pool) + .await?; + + let day_of_week: Vec<(String, i64)> = sqlx::query_as(&day_of_week) + .bind(&h.command) + .fetch_all(&self.pool) + .await?; + + let duration_over_time: Vec<(String, f64)> = sqlx::query_as(&duration_over_time) + .bind(&h.command) + .fetch_all(&self.pool) + .await?; + + let duration_over_time = duration_over_time + .iter() + .map(|f| (f.0.clone(), f.1.round() as i64)) + .collect(); + + Ok(HistoryStats { + next, + previous: prev, + total: total.0 as u64, + average_duration: average.0 as u64, + exits, + day_of_week, + duration_over_time, + }) + } + + pub(crate) async fn get_dups(&self, before: i64, dupkeep: u32) -> Result<Vec<History>> { + let res = sqlx::query( + "SELECT * FROM ( + SELECT *, ROW_NUMBER() + OVER (PARTITION BY command, cwd, hostname ORDER BY timestamp DESC) + AS rn + FROM history + ) sub + WHERE rn > ?1 and timestamp < ?2; + ", + ) + .bind(dupkeep) + .bind(before) + .map(Self::query_history_inner) + .fetch_all(&self.pool) + .await?; + + Ok(res) + } +} + +pub(crate) struct Paged { + database: ClientSqlite, + page_size: usize, + last_id: Option<String>, + include_deleted: bool, + unique: bool, +} + +impl Paged { + pub(crate) fn new( + database: ClientSqlite, + page_size: usize, + include_deleted: bool, + unique: bool, + ) -> Self { + Self { + database, + page_size, + last_id: None, + include_deleted, + unique, + } + } + + pub(crate) async fn next(&mut self) -> Result<Option<Vec<History>>> { + let mut query = SqlBuilder::select_from(SqlName::new("history").alias("h").baquoted()); + + query.field("*").order_desc("id"); + + if !self.include_deleted { + query.and_where_is_null("deleted_at"); + } + + if self.unique { + // We want to deduplicate on command, but the user can search via cwd, hostname, and session. + // Without those fields, filter modes won't work right. With those fields, we get duplicates. + // This must be handled upstream. + query + .group_by("command, cwd, hostname, session") + .having("max(timestamp)"); + } + + query.limit(self.page_size); + + if let Some(last_id) = &self.last_id { + query.and_where_lt("id", quote(last_id)); + } + + let query = query.sql().expect("bug in list query. please report"); + let res = self.database.query_history(&query).await?; + + if res.is_empty() { + Ok(None) + } else { + self.last_id = Some(res.last().unwrap().id.0.clone()); + Ok(Some(res)) + } + } +} + +trait SqlBuilderExt { + fn fuzzy_condition<S: ToString, T: ToString>( + &mut self, + field: S, + mask: T, + inverse: bool, + glob: bool, + is_or: bool, + ) -> &mut Self; +} + +impl SqlBuilderExt for SqlBuilder { + /// adapted from the sql-builder *like functions + fn fuzzy_condition<S: ToString, T: ToString>( + &mut self, + field: S, + mask: T, + inverse: bool, + glob: bool, + is_or: bool, + ) -> &mut Self { + let mut cond = field.to_string(); + if inverse { + cond.push_str(" NOT"); + } + if glob { + cond.push_str(" GLOB '"); + } else { + cond.push_str(" LIKE '"); + } + cond.push_str(&esc(mask.to_string())); + cond.push('\''); + if is_or { + self.or_where(cond) + } else { + self.and_where(cond) + } + } +} + +#[cfg(test)] +mod test { + use crate::aclient::settings::test_local_timeout; + + use super::{ + ClientSqlite, Context, FilterMode, History, OffsetDateTime, OptFilters, Result, SearchMode, + }; + use std::time::{Duration, Instant}; + + async fn assert_search_eq( + db: &ClientSqlite, + mode: SearchMode, + filter_mode: FilterMode, + query: &str, + expected: usize, + ) -> Result<Vec<History>> { + let context = Context { + hostname: "test:host".to_string(), + session: "beepboopiamasession".to_string(), + cwd: "/home/ellie".to_string(), + host_id: "test-host".to_string(), + git_root: None, + }; + + let results = db + .search( + mode, + filter_mode, + &context, + query, + OptFilters { + ..Default::default() + }, + ) + .await?; + + assert_eq!( + results.len(), + expected, + "query \"{}\", commands: {:?}", + query, + results.iter().map(|a| &a.command).collect::<Vec<&String>>() + ); + Ok(results) + } + + async fn assert_search_commands( + db: &ClientSqlite, + mode: SearchMode, + filter_mode: FilterMode, + query: &str, + expected_commands: Vec<&str>, + ) { + let results = assert_search_eq(db, mode, filter_mode, query, expected_commands.len()) + .await + .unwrap(); + let commands: Vec<&str> = results.iter().map(|a| a.command.as_str()).collect(); + assert_eq!(commands, expected_commands); + } + + async fn new_history_item(db: &mut ClientSqlite, cmd: &str) -> Result<()> { + let mut captured: History = History::capture() + .timestamp(OffsetDateTime::now_utc()) + .command(cmd) + .cwd("/home/ellie") + .build() + .into(); + + captured.exit = 0; + captured.duration = 1; + captured.session = "beep boop".to_string(); + captured.hostname = "booop".to_string(); + + db.save(&captured).await + } + + #[tokio::test(flavor = "multi_thread")] + async fn test_search_prefix() { + let mut db = ClientSqlite::new("sqlite::memory:", test_local_timeout()) + .await + .unwrap(); + new_history_item(&mut db, "ls /home/ellie").await.unwrap(); + + assert_search_eq(&db, SearchMode::Prefix, FilterMode::Global, "ls", 1) + .await + .unwrap(); + assert_search_eq(&db, SearchMode::Prefix, FilterMode::Global, "/home", 0) + .await + .unwrap(); + assert_search_eq(&db, SearchMode::Prefix, FilterMode::Global, "ls ", 0) + .await + .unwrap(); + } + + #[tokio::test(flavor = "multi_thread")] + async fn test_search_fulltext() { + let mut db = ClientSqlite::new("sqlite::memory:", test_local_timeout()) + .await + .unwrap(); + new_history_item(&mut db, "ls /home/ellie").await.unwrap(); + + assert_search_eq(&db, SearchMode::FullText, FilterMode::Global, "ls", 1) + .await + .unwrap(); + assert_search_eq(&db, SearchMode::FullText, FilterMode::Global, "/home", 1) + .await + .unwrap(); + assert_search_eq(&db, SearchMode::FullText, FilterMode::Global, "ls ho", 1) + .await + .unwrap(); + assert_search_eq(&db, SearchMode::FullText, FilterMode::Global, "hm", 0) + .await + .unwrap(); + + // regex + assert_search_eq(&db, SearchMode::FullText, FilterMode::Global, "r/^ls ", 1) + .await + .unwrap(); + assert_search_eq( + &db, + SearchMode::FullText, + FilterMode::Global, + "r/ls / ie$", + 1, + ) + .await + .unwrap(); + assert_search_eq( + &db, + SearchMode::FullText, + FilterMode::Global, + "r/ls / !ie", + 0, + ) + .await + .unwrap(); + assert_search_eq( + &db, + SearchMode::FullText, + FilterMode::Global, + "meow r/ls/", + 0, + ) + .await + .unwrap(); + assert_search_eq(&db, SearchMode::FullText, FilterMode::Global, "r//hom/", 1) + .await + .unwrap(); + assert_search_eq( + &db, + SearchMode::FullText, + FilterMode::Global, + "r//home//", + 1, + ) + .await + .unwrap(); + assert_search_eq( + &db, + SearchMode::FullText, + FilterMode::Global, + "r//home///", + 0, + ) + .await + .unwrap(); + assert_search_eq(&db, SearchMode::FullText, FilterMode::Global, "/home.*e", 0) + .await + .unwrap(); + assert_search_eq( + &db, + SearchMode::FullText, + FilterMode::Global, + "r/home.*e", + 1, + ) + .await + .unwrap(); + } + + #[tokio::test(flavor = "multi_thread")] + async fn test_search_fuzzy() { + let mut db = ClientSqlite::new("sqlite::memory:", test_local_timeout()) + .await + .unwrap(); + new_history_item(&mut db, "ls /home/ellie").await.unwrap(); + new_history_item(&mut db, "ls /home/frank").await.unwrap(); + new_history_item(&mut db, "cd /home/Ellie").await.unwrap(); + new_history_item(&mut db, "/home/ellie/.bin/rustup") + .await + .unwrap(); + + assert_search_eq(&db, SearchMode::Fuzzy, FilterMode::Global, "ls /", 3) + .await + .unwrap(); + assert_search_eq(&db, SearchMode::Fuzzy, FilterMode::Global, "ls/", 2) + .await + .unwrap(); + assert_search_eq(&db, SearchMode::Fuzzy, FilterMode::Global, "l/h/", 2) + .await + .unwrap(); + assert_search_eq(&db, SearchMode::Fuzzy, FilterMode::Global, "/h/e", 3) + .await + .unwrap(); + assert_search_eq(&db, SearchMode::Fuzzy, FilterMode::Global, "/hmoe/", 0) + .await + .unwrap(); + assert_search_eq(&db, SearchMode::Fuzzy, FilterMode::Global, "ellie/home", 0) + .await + .unwrap(); + assert_search_eq(&db, SearchMode::Fuzzy, FilterMode::Global, "lsellie", 1) + .await + .unwrap(); + assert_search_eq(&db, SearchMode::Fuzzy, FilterMode::Global, " ", 4) + .await + .unwrap(); + + // single term operators + assert_search_eq(&db, SearchMode::Fuzzy, FilterMode::Global, "^ls", 2) + .await + .unwrap(); + assert_search_eq(&db, SearchMode::Fuzzy, FilterMode::Global, "'ls", 2) + .await + .unwrap(); + assert_search_eq(&db, SearchMode::Fuzzy, FilterMode::Global, "ellie$", 2) + .await + .unwrap(); + assert_search_eq(&db, SearchMode::Fuzzy, FilterMode::Global, "!^ls", 2) + .await + .unwrap(); + assert_search_eq(&db, SearchMode::Fuzzy, FilterMode::Global, "!ellie", 1) + .await + .unwrap(); + assert_search_eq(&db, SearchMode::Fuzzy, FilterMode::Global, "!ellie$", 2) + .await + .unwrap(); + + // multiple terms + assert_search_eq(&db, SearchMode::Fuzzy, FilterMode::Global, "ls !ellie", 1) + .await + .unwrap(); + assert_search_eq(&db, SearchMode::Fuzzy, FilterMode::Global, "^ls !e$", 1) + .await + .unwrap(); + assert_search_eq(&db, SearchMode::Fuzzy, FilterMode::Global, "home !^ls", 2) + .await + .unwrap(); + assert_search_eq( + &db, + SearchMode::Fuzzy, + FilterMode::Global, + "'frank | 'rustup", + 2, + ) + .await + .unwrap(); + assert_search_eq( + &db, + SearchMode::Fuzzy, + FilterMode::Global, + "'frank | 'rustup 'ls", + 1, + ) + .await + .unwrap(); + + // case matching + assert_search_eq(&db, SearchMode::Fuzzy, FilterMode::Global, "Ellie", 1) + .await + .unwrap(); + + // regex + assert_search_eq(&db, SearchMode::Fuzzy, FilterMode::Global, "r/^ls ", 2) + .await + .unwrap(); + assert_search_eq(&db, SearchMode::Fuzzy, FilterMode::Global, "r/[Ee]llie", 3) + .await + .unwrap(); + assert_search_eq(&db, SearchMode::Fuzzy, FilterMode::Global, "/h/e r/^ls ", 1) + .await + .unwrap(); + } + + #[tokio::test(flavor = "multi_thread")] + async fn test_search_reordered_fuzzy() { + let mut db = ClientSqlite::new("sqlite::memory:", test_local_timeout()) + .await + .unwrap(); + // test ordering of results: we should choose the first, even though it happened longer ago. + + new_history_item(&mut db, "curl").await.unwrap(); + new_history_item(&mut db, "corburl").await.unwrap(); + + // if fuzzy reordering is on, it should come back in a more sensible order + assert_search_commands( + &db, + SearchMode::Fuzzy, + FilterMode::Global, + "curl", + vec!["curl", "corburl"], + ) + .await; + + assert_search_eq(&db, SearchMode::Fuzzy, FilterMode::Global, "xxxx", 0) + .await + .unwrap(); + assert_search_eq(&db, SearchMode::Fuzzy, FilterMode::Global, "", 2) + .await + .unwrap(); + } + + #[tokio::test(flavor = "multi_thread")] + #[expect(clippy::similar_names)] + async fn test_paged_basic() { + let mut db = ClientSqlite::new("sqlite::memory:", test_local_timeout()) + .await + .unwrap(); + + // Add 5 history items + for i in 0..5 { + new_history_item(&mut db, &format!("command{}", i)) + .await + .unwrap(); + } + + // Create a paged iterator with page_size of 2 + let mut paged = db.all_paged(2, false, false); + + // First page should have 2 items + let page1 = paged.next().await.unwrap(); + assert!(page1.is_some()); + assert_eq!(page1.unwrap().len(), 2); + + // Second page should have 2 items + let page2 = paged.next().await.unwrap(); + assert!(page2.is_some()); + assert_eq!(page2.unwrap().len(), 2); + + // Third page should have 1 item + let page3 = paged.next().await.unwrap(); + assert!(page3.is_some()); + assert_eq!(page3.unwrap().len(), 1); + + // Fourth page should be None (exhausted) + let page4 = paged.next().await.unwrap(); + assert!(page4.is_none()); + } + + #[tokio::test(flavor = "multi_thread")] + async fn test_paged_empty() { + let db = ClientSqlite::new("sqlite::memory:", test_local_timeout()) + .await + .unwrap(); + + // Create a paged iterator on empty database + let mut paged = db.all_paged(10, false, false); + + // Should return None immediately + let page = paged.next().await.unwrap(); + assert!(page.is_none()); + } + + #[tokio::test(flavor = "multi_thread")] + async fn test_paged_unique() { + let mut db = ClientSqlite::new("sqlite::memory:", test_local_timeout()) + .await + .unwrap(); + + // Add duplicate commands + new_history_item(&mut db, "duplicate").await.unwrap(); + new_history_item(&mut db, "duplicate").await.unwrap(); + new_history_item(&mut db, "unique1").await.unwrap(); + new_history_item(&mut db, "unique2").await.unwrap(); + + // Without unique flag - should get all 4 + let mut paged = db.all_paged(10, false, false); + let page = paged.next().await.unwrap().unwrap(); + assert_eq!(page.len(), 4); + + // With unique flag - should get 3 (duplicates collapsed) + let mut paged_unique = db.all_paged(10, false, true); + let paged_unique = paged_unique.next().await.unwrap().unwrap(); + assert_eq!(paged_unique.len(), 3); + } + + #[tokio::test(flavor = "multi_thread")] + async fn test_search_bench_dupes() { + let context = Context { + hostname: "test:host".to_string(), + session: "beepboopiamasession".to_string(), + cwd: "/home/ellie".to_string(), + host_id: "test-host".to_string(), + git_root: None, + }; + + let mut db = ClientSqlite::new("sqlite::memory:", test_local_timeout()) + .await + .unwrap(); + for _i in 1..10000 { + new_history_item(&mut db, "i am a duplicated command") + .await + .unwrap(); + } + let start = Instant::now(); + let _results = db + .search( + SearchMode::Fuzzy, + FilterMode::Global, + &context, + "", + OptFilters { + ..Default::default() + }, + ) + .await + .unwrap(); + let duration = start.elapsed(); + + assert!(duration < Duration::from_secs(15)); + } +} + +pub(crate) struct QueryTokenizer<'a> { + query: &'a str, + last_pos: usize, +} + +pub(crate) enum QueryToken<'a> { + Match(&'a str, bool), + MatchStart(&'a str, bool), + MatchEnd(&'a str, bool), + MatchFull(&'a str, bool), + Or, + Regex(&'a str), +} + +impl QueryToken<'_> { + pub(crate) fn has_uppercase(&self) -> bool { + match self { + Self::Match(term, _) + | Self::MatchStart(term, _) + | Self::MatchEnd(term, _) + | Self::MatchFull(term, _) => term.contains(char::is_uppercase), + _ => false, + } + } + + pub(crate) fn is_inverse(&self) -> bool { + match self { + Self::Match(_, inv) + | Self::MatchStart(_, inv) + | Self::MatchEnd(_, inv) + | Self::MatchFull(_, inv) => *inv, + _ => false, + } + } +} + +impl<'a> QueryTokenizer<'a> { + pub(crate) fn new(query: &'a str) -> Self { + Self { query, last_pos: 0 } + } +} + +impl<'a> Iterator for QueryTokenizer<'a> { + type Item = QueryToken<'a>; + fn next(&mut self) -> Option<Self::Item> { + let remaining = &self.query[self.last_pos..]; + if remaining.is_empty() { + return None; + } + + if let Some(remaining) = remaining.strip_prefix("r/") { + let (regex, next_pos) = if let Some(end) = remaining.find("/ ") { + (&remaining[..end], self.last_pos + 2 + end + 2) + } else if let Some(remaining) = remaining.strip_suffix('/') { + (remaining, self.query.len()) + } else { + (remaining, self.query.len()) + }; + self.last_pos = next_pos; + Some(QueryToken::Regex(regex)) + } else { + let (mut part, next_pos) = if let Some(sp) = remaining.find(' ') { + (&remaining[..sp], self.last_pos + sp + 1) + } else { + (remaining, self.query.len()) + }; + self.last_pos = next_pos; + + if part == "|" { + return Some(QueryToken::Or); + } + + let is_inverse = part.strip_prefix('!').is_some_and(|s| { + part = s; + true + }); + + #[expect(clippy::option_if_let_else, reason = "It's too ugly")] + let token = if let Some(s) = part.strip_prefix('^') { + QueryToken::MatchStart(s, is_inverse) + } else if let Some(s) = part.strip_suffix('$') { + QueryToken::MatchEnd(s, is_inverse) + } else if let Some(s) = part.strip_prefix('\'') { + QueryToken::MatchFull(s, is_inverse) + } else { + QueryToken::Match(part, is_inverse) + }; + Some(token) + } + } +} diff --git a/crates/daemon/src/aclient/encryption.rs b/crates/daemon/src/aclient/encryption.rs new file mode 100644 index 00000000..220ac74e --- /dev/null +++ b/crates/daemon/src/aclient/encryption.rs @@ -0,0 +1,142 @@ +// The general idea is that we NEVER send cleartext history to the server +// This way the odds of anything private ending up where it should not are +// very low +// The server authenticates via the usual username and password. This has +// nothing to do with the encryption, and is purely authentication! The client +// generates its own secret key, and encrypts all shell history with libsodium's +// secretbox. The data is then sent to the server, where it is stored. All +// clients must share the secret in order to be able to sync, as it is needed +// to decrypt + +use std::io::prelude::Write; + +use base64::prelude::{BASE64_STANDARD, Engine}; +pub(crate) use crypto_secretbox::Key; +use crypto_secretbox::{KeyInit, XSalsa20Poly1305, aead::OsRng}; +use eyre::{Context, Result, bail, ensure, eyre}; +use fs_err as fs; +use rmp::Marker; + +use crate::aclient::settings::Settings; + +pub(crate) fn generate_encoded_key() -> Result<(Key, String)> { + let key = XSalsa20Poly1305::generate_key(&mut OsRng); + let encoded = encode_key(&key)?; + + Ok((key, encoded)) +} + +pub(crate) fn new_key(settings: &Settings) -> Result<Key> { + if settings.sync.encryption_key()?.is_some() { + bail!("key already exists! cannot overwrite"); + } else if let Some(path) = settings.sync.encryption_key_path.as_ref() { + let (key, encoded) = generate_encoded_key()?; + + let mut file = fs::File::create(path)?; + file.write_all(encoded.as_bytes())?; + + Ok(key) + } else { + bail!("No key-path set, cannot generate key") + } +} + +// Loads the secret key, will create + save if it doesn't exist +pub(crate) fn load_key(settings: &Settings) -> Result<Key> { + if let Some(key) = settings.sync.encryption_key()? { + Ok(key) + } else { + Ok(new_key(settings)?) + } +} + +pub(crate) fn encode_key(key: &Key) -> Result<String> { + let mut buf = vec![]; + rmp::encode::write_array_len(&mut buf, key.len() as u32) + .wrap_err("could not encode key to message pack")?; + for b in key { + rmp::encode::write_uint(&mut buf, u64::from(*b)) + .wrap_err("could not encode key to message pack")?; + } + let buf = BASE64_STANDARD.encode(buf); + + Ok(buf) +} + +pub(crate) fn decode_key(key: &str) -> Result<Key> { + use rmp::decode; + + let buf = BASE64_STANDARD + .decode(key.trim_end()) + .wrap_err("encryption key is not a valid base64 encoding")?; + + // old code wrote the key as a fixed length array of 32 bytes + // new code writes the key with a length prefix + if let Ok(key) = <[u8; 32]>::try_from(&*buf) { + Ok(key.into()) + } else { + let mut bytes = decode::Bytes::new(&buf); + + match Marker::from_u8(buf[0]) { + Marker::Bin8 => { + let len = decode::read_bin_len(&mut bytes).map_err(|err| eyre!("{err:?}"))?; + ensure!(len == 32, "encryption key is not the correct size"); + let key = <[u8; 32]>::try_from(bytes.remaining_slice()) + .context("could not decode encryption key")?; + Ok(key.into()) + } + Marker::Array16 => { + let len = decode::read_array_len(&mut bytes).map_err(|err| eyre!("{err:?}"))?; + ensure!(len == 32, "encryption key is not the correct size"); + + let mut key = Key::default(); + for i in &mut key { + *i = decode::read_int(&mut bytes).map_err(|err| eyre!("{err:?}"))?; + } + Ok(key) + } + _ => bail!("could not decode encryption key"), + } + } +} + +#[cfg(test)] +mod test { + #[test] + fn key_encodings() { + use super::{Key, decode_key, encode_key}; + + // a history of our key encodings. + // v11.0.0 xCAbWypb0msJ2Kq+8j4GVEWUlDX7deKnrTRSIopuqXxc5Q== + // v12.0.0 xCAbWypb0msJ2Kq+8j4GVEWUlDX7deKnrTRSIopuqXxc5Q== + // v13.0.0 xCAbWypb0msJ2Kq+8j4GVEWUlDX7deKnrTRSIopuqXxc5Q== + // v13.0.1 xCAbWypb0msJ2Kq+8j4GVEWUlDX7deKnrTRSIopuqXxc5Q== + // v14.0.0 xCAbWypb0msJ2Kq+8j4GVEWUlDX7deKnrTRSIopuqXxc5Q== + // v14.0.1 xCAbWypb0msJ2Kq+8j4GVEWUlDX7deKnrTRSIopuqXxc5Q== + // c7d89c1 3AAgG1sqW8zSawnM2MyqzL7M8j4GVEXMlMyUNcz7dczizKfMrTRSIsyKbsypfFzM5Q== (https://github.com/ellie/atuin/pull/805) + // b53ca35 3AAgG1sqW8zSawnM2MyqzL7M8j4GVEXMlMyUNcz7dczizKfMrTRSIsyKbsypfFzM5Q== (https://github.com/ellie/atuin/pull/974) + // v15.0.0 3AAgG1sqW8zSawnM2MyqzL7M8j4GVEXMlMyUNcz7dczizKfMrTRSIsyKbsypfFzM5Q== + // b8b57c8 xCAbWypb0msJ2Kq+8j4GVEWUlDX7deKnrTRSIopuqXxc5Q== (https://github.com/ellie/atuin/pull/1057) + // 8c94d79 3AAgG1sqW8zSawnM2MyqzL7M8j4GVEXMlMyUNcz7dczizKfMrTRSIsyKbsypfFzM5Q== (https://github.com/ellie/atuin/pull/1089) + + let key = Key::from([ + 27, 91, 42, 91, 210, 107, 9, 216, 170, 190, 242, 62, 6, 84, 69, 148, 148, 53, 251, 117, + 226, 167, 173, 52, 82, 34, 138, 110, 169, 124, 92, 229, + ]); + + assert_eq!( + encode_key(&key).unwrap(), + "3AAgG1sqW8zSawnM2MyqzL7M8j4GVEXMlMyUNcz7dczizKfMrTRSIsyKbsypfFzM5Q==" + ); + + // key encodings we have to support + let valid_encodings = [ + "xCAbWypb0msJ2Kq+8j4GVEWUlDX7deKnrTRSIopuqXxc5Q==", + "3AAgG1sqW8zSawnM2MyqzL7M8j4GVEXMlMyUNcz7dczizKfMrTRSIsyKbsypfFzM5Q==", + ]; + + for k in valid_encodings { + assert_eq!(decode_key(k).expect(k), key); + } + } +} diff --git a/crates/daemon/src/aclient/history/builder.rs b/crates/daemon/src/aclient/history/builder.rs new file mode 100644 index 00000000..daa4ef49 --- /dev/null +++ b/crates/daemon/src/aclient/history/builder.rs @@ -0,0 +1,154 @@ +use typed_builder::TypedBuilder; + +use super::History; + +/// Builder for a history entry that is imported from shell history. +/// +/// The only two required fields are `timestamp` and `command`. +#[derive(Debug, Clone, TypedBuilder)] +pub(crate) struct HistoryImported { + timestamp: time::OffsetDateTime, + #[builder(setter(into))] + command: String, + #[builder(default = "unknown".into(), setter(into))] + cwd: String, + #[builder(default = -1)] + exit: i64, + #[builder(default = -1)] + duration: i64, + #[builder(default, setter(strip_option, into))] + session: Option<String>, + #[builder(default, setter(strip_option, into))] + hostname: Option<String>, + #[builder(default, setter(strip_option, into))] + author: Option<String>, + #[builder(default, setter(strip_option, into))] + intent: Option<String>, +} + +impl From<HistoryImported> for History { + fn from(imported: HistoryImported) -> Self { + Self::new( + imported.timestamp, + imported.command, + imported.cwd, + imported.exit, + imported.duration, + imported.session, + imported.hostname, + imported.author, + imported.intent, + None, + ) + } +} + +/// Builder for a history entry that is captured via hook. +/// +/// This builder is used only at the `start` step of the hook, +/// so it doesn't have any fields which are known only after +/// the command is finished, such as `exit` or `duration`. +#[derive(Debug, Clone, TypedBuilder)] +pub(crate) struct HistoryCaptured { + timestamp: time::OffsetDateTime, + #[builder(setter(into))] + command: String, + #[builder(setter(into))] + cwd: String, + #[builder(default, setter(strip_option, into))] + author: Option<String>, + #[builder(default, setter(strip_option, into))] + intent: Option<String>, +} + +impl From<HistoryCaptured> for History { + fn from(captured: HistoryCaptured) -> Self { + Self::new( + captured.timestamp, + captured.command, + captured.cwd, + -1, + -1, + None, + None, + captured.author, + captured.intent, + None, + ) + } +} + +/// Builder for a history entry that is loaded from the database. +/// +/// All fields are required, as they are all present in the database. +#[derive(Debug, Clone, TypedBuilder)] +pub(crate) struct HistoryFromDb { + id: String, + timestamp: time::OffsetDateTime, + command: String, + cwd: String, + exit: i64, + duration: i64, + session: String, + hostname: String, + author: String, + intent: Option<String>, + deleted_at: Option<time::OffsetDateTime>, +} + +impl From<HistoryFromDb> for History { + fn from(from_db: HistoryFromDb) -> Self { + Self { + id: from_db.id.into(), + timestamp: from_db.timestamp, + exit: from_db.exit, + command: from_db.command, + cwd: from_db.cwd, + duration: from_db.duration, + session: from_db.session, + hostname: from_db.hostname, + author: from_db.author, + intent: from_db.intent, + deleted_at: from_db.deleted_at, + } + } +} + +/// Builder for a history entry that is captured via hook and sent to the daemon +/// +/// This builder is similar to Capture, but we just require more information up front. +/// For the old setup, we could just rely on `History::new` to read some of the missing +/// data. This is no longer the case. +#[derive(Debug, Clone, TypedBuilder)] +pub(crate) struct HistoryDaemonCapture { + timestamp: time::OffsetDateTime, + #[builder(setter(into))] + command: String, + #[builder(setter(into))] + cwd: String, + #[builder(setter(into))] + session: String, + #[builder(setter(into))] + hostname: String, + #[builder(default, setter(strip_option, into))] + author: Option<String>, + #[builder(default, setter(strip_option, into))] + intent: Option<String>, +} + +impl From<HistoryDaemonCapture> for History { + fn from(captured: HistoryDaemonCapture) -> Self { + Self::new( + captured.timestamp, + captured.command, + captured.cwd, + -1, + -1, + Some(captured.session), + Some(captured.hostname), + captured.author, + captured.intent, + None, + ) + } +} diff --git a/crates/daemon/src/aclient/history/mod.rs b/crates/daemon/src/aclient/history/mod.rs new file mode 100644 index 00000000..09d24169 --- /dev/null +++ b/crates/daemon/src/aclient/history/mod.rs @@ -0,0 +1,686 @@ +use core::fmt::Formatter; +use rmp::decode::DecodeStringError; +use rmp::decode::ValueReadError; +use rmp::{Marker, decode::Bytes}; +use std::env; +use std::fmt::Display; + +use turtle_common::record::DecryptedData; +use turtle_common::utils::uuid_v7; + +use eyre::{Result, bail, eyre}; + +use crate::aclient::secrets::SECRET_PATTERNS_RE; +use crate::aclient::settings::Settings; +use crate::aclient::utils::get_host_user; +use time::OffsetDateTime; + +mod builder; +pub(crate) mod store; + +pub(crate) const HISTORY_VERSION_V0: &str = "v0"; +pub(crate) const HISTORY_VERSION_V1: &str = "v1"; +const HISTORY_RECORD_VERSION_V0: u16 = 0; +const HISTORY_RECORD_VERSION_V1: u16 = 1; +pub(crate) const HISTORY_VERSION: &str = HISTORY_VERSION_V1; +pub(crate) const HISTORY_TAG: &str = "history"; +const HISTORY_AUTHOR_ENV: &str = "ATUIN_HISTORY_AUTHOR"; +const HISTORY_INTENT_ENV: &str = "ATUIN_HISTORY_INTENT"; + +#[derive(Clone, Debug, Eq, PartialEq, Hash)] +pub struct HistoryId(pub(crate) String); + +impl Display for HistoryId { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!(f, "{}", self.0) + } +} + +impl From<String> for HistoryId { + fn from(s: String) -> Self { + Self(s) + } +} + +/// Client-side history entry. +/// +/// Client stores data unencrypted, and only encrypts it before sending to the server. +/// +/// To create a new history entry, use one of the builders: +/// - [`History::import()`] to import an entry from the shell history file +/// - [`History::capture()`] to capture an entry via hook +/// - [`History::from_db()`] to create an instance from the database entry +// +// ## Implementation Notes +// +// New fields must be added to `History::{serialize,deserialize}` in a backwards +// compatible way (sensible defaults and careful `nfields` handling). +#[derive(Debug, Clone, PartialEq, Eq, sqlx::FromRow)] +pub struct History { + /// A client-generated ID, used to identify the entry when syncing. + /// + /// Stored as `client_id` in the database. + pub(crate) id: HistoryId, + + /// When the command was run. + pub(crate) timestamp: OffsetDateTime, + + /// How long the command took to run. + pub(crate) duration: i64, + + /// The exit code of the command. + pub(crate) exit: i64, + + /// The command that was run. + pub(crate) command: String, + + /// The current working directory when the command was run. + pub(crate) cwd: String, + + /// The session ID, associated with a terminal session. + pub(crate) session: String, + + /// The hostname of the machine the command was run on. + pub(crate) hostname: String, + + /// Who wrote this command (human user or automation/agent identity). + pub(crate) author: String, + + /// Optional rationale for why the command was executed. + pub(crate) intent: Option<String>, + + /// Timestamp, which is set when the entry is deleted, allowing a soft delete. + pub(crate) deleted_at: Option<OffsetDateTime>, +} + +#[derive(Debug, Clone, PartialEq, Eq, sqlx::FromRow)] +pub(crate) struct HistoryStats { + /// The command that was ran after this one in the session + pub(crate) next: Option<History>, + + /// The command that was ran before this one in the session + pub(crate) previous: Option<History>, + + /// How many times has this command been ran? + pub(crate) total: u64, + + pub(crate) average_duration: u64, + + pub(crate) exits: Vec<(i64, i64)>, + + pub(crate) day_of_week: Vec<(String, i64)>, + + pub(crate) duration_over_time: Vec<(String, i64)>, +} + +impl History { + pub(crate) fn author_from_hostname(hostname: &str) -> String { + hostname + .split_once(':') + .map_or_else(|| hostname.to_owned(), |(_, user)| user.to_owned()) + } + + fn normalize_optional_field(field: Option<String>) -> Option<String> { + field.and_then(|value| { + let trimmed = value.trim(); + if trimmed.is_empty() { + None + } else { + Some(trimmed.to_owned()) + } + }) + } + + #[expect(clippy::too_many_arguments)] + fn new( + timestamp: OffsetDateTime, + command: String, + cwd: String, + exit: i64, + duration: i64, + session: Option<String>, + hostname: Option<String>, + author: Option<String>, + intent: Option<String>, + deleted_at: Option<OffsetDateTime>, + ) -> Self { + let session = session + .or_else(|| env::var("ATUIN_SESSION").ok()) + .unwrap_or_else(|| uuid_v7().as_simple().to_string()); + let hostname = hostname.unwrap_or_else(get_host_user); + let author = Self::normalize_optional_field(author) + .or_else(|| Self::normalize_optional_field(env::var(HISTORY_AUTHOR_ENV).ok())) + .unwrap_or_else(|| Self::author_from_hostname(hostname.as_str())); + let intent = Self::normalize_optional_field(intent) + .or_else(|| Self::normalize_optional_field(env::var(HISTORY_INTENT_ENV).ok())); + + Self { + id: uuid_v7().as_simple().to_string().into(), + timestamp, + command, + cwd, + exit, + duration, + session, + hostname, + author, + intent, + deleted_at, + } + } + + pub(crate) fn serialize(&self) -> Result<DecryptedData> { + // This is pretty much the same as what we used for the old history, with one difference - + // it uses integers for timestamps rather than a string format. + + use rmp::encode; + + let mut output = vec![]; + + // write the version + encode::write_u16(&mut output, HISTORY_RECORD_VERSION_V1)?; + let include_intent = self.intent.is_some(); + encode::write_array_len(&mut output, 10 + u32::from(include_intent))?; + + encode::write_str(&mut output, &self.id.0)?; + encode::write_u64(&mut output, self.timestamp.unix_timestamp_nanos() as u64)?; + encode::write_sint(&mut output, self.duration)?; + encode::write_sint(&mut output, self.exit)?; + encode::write_str(&mut output, &self.command)?; + encode::write_str(&mut output, &self.cwd)?; + encode::write_str(&mut output, &self.session)?; + encode::write_str(&mut output, &self.hostname)?; + + match self.deleted_at { + Some(d) => encode::write_u64(&mut output, d.unix_timestamp_nanos() as u64)?, + None => encode::write_nil(&mut output)?, + } + + encode::write_str(&mut output, self.author.as_str())?; + if let Some(intent) = &self.intent { + encode::write_str(&mut output, intent.as_str())?; + } + + Ok(DecryptedData(output)) + } + + fn read_optional_string(bytes: &[u8]) -> Result<(Option<String>, &[u8])> { + use rmp::decode; + + fn error_report<E: std::fmt::Debug>(err: E) -> eyre::Report { + eyre!("{err:?}") + } + + match decode::read_str_from_slice(bytes) { + Ok((value, bytes)) => Ok((Some(value.to_owned()), bytes)), + Err(DecodeStringError::TypeMismatch(Marker::Null)) => { + let mut cursor = Bytes::new(bytes); + decode::read_nil(&mut cursor).map_err(error_report)?; + + Ok((None, cursor.remaining_slice())) + } + Err(err) => Err(error_report(err)), + } + } + + fn deserialize_v0(bytes: &[u8]) -> Result<Self> { + use rmp::decode; + + fn error_report<E: std::fmt::Debug>(err: E) -> eyre::Report { + eyre!("{err:?}") + } + + let mut bytes = Bytes::new(bytes); + + let version = decode::read_u16(&mut bytes).map_err(error_report)?; + + if version != HISTORY_RECORD_VERSION_V0 { + bail!("expected decoding v0 record, found v{version}"); + } + + let nfields = decode::read_array_len(&mut bytes).map_err(error_report)?; + + if nfields != 9 { + bail!("cannot decrypt history from a different version of Atuin"); + } + + let bytes = bytes.remaining_slice(); + let (id, bytes) = decode::read_str_from_slice(bytes).map_err(error_report)?; + + let mut bytes = Bytes::new(bytes); + let timestamp = decode::read_u64(&mut bytes).map_err(error_report)?; + let duration = decode::read_int(&mut bytes).map_err(error_report)?; + let exit = decode::read_int(&mut bytes).map_err(error_report)?; + + let bytes = bytes.remaining_slice(); + let (command, bytes) = decode::read_str_from_slice(bytes).map_err(error_report)?; + let (cwd, bytes) = decode::read_str_from_slice(bytes).map_err(error_report)?; + let (session, bytes) = decode::read_str_from_slice(bytes).map_err(error_report)?; + let (hostname, bytes) = decode::read_str_from_slice(bytes).map_err(error_report)?; + + let mut bytes = Bytes::new(bytes); + + let (deleted_at, bytes) = match decode::read_u64(&mut bytes) { + Ok(unix) => (Some(unix), bytes.remaining_slice()), + // we accept null here + Err(ValueReadError::TypeMismatch(Marker::Null)) => (None, bytes.remaining_slice()), + Err(err) => return Err(error_report(err)), + }; + if !bytes.is_empty() { + bail!("trailing bytes in encoded history. malformed") + } + + Ok(Self { + id: id.to_owned().into(), + timestamp: OffsetDateTime::from_unix_timestamp_nanos(i128::from(timestamp))?, + duration, + exit, + command: command.to_owned(), + cwd: cwd.to_owned(), + session: session.to_owned(), + hostname: hostname.to_owned(), + author: Self::author_from_hostname(hostname), + intent: None, + deleted_at: deleted_at + .map(|t| OffsetDateTime::from_unix_timestamp_nanos(i128::from(t))) + .transpose()?, + }) + } + + fn deserialize_v1(bytes: &[u8]) -> Result<Self> { + use rmp::decode; + + fn error_report<E: std::fmt::Debug>(err: E) -> eyre::Report { + eyre!("{err:?}") + } + + let mut bytes = Bytes::new(bytes); + + let version = decode::read_u16(&mut bytes).map_err(error_report)?; + + if version != HISTORY_RECORD_VERSION_V1 { + bail!("expected decoding v1 record, found v{version}"); + } + + let nfields = decode::read_array_len(&mut bytes).map_err(error_report)?; + + if !(10..=11).contains(&nfields) { + bail!("cannot decrypt history from a different version of Atuin"); + } + + let bytes = bytes.remaining_slice(); + let (id, bytes) = decode::read_str_from_slice(bytes).map_err(error_report)?; + + let mut bytes = Bytes::new(bytes); + let timestamp = decode::read_u64(&mut bytes).map_err(error_report)?; + let duration = decode::read_int(&mut bytes).map_err(error_report)?; + let exit = decode::read_int(&mut bytes).map_err(error_report)?; + + let bytes = bytes.remaining_slice(); + let (command, bytes) = decode::read_str_from_slice(bytes).map_err(error_report)?; + let (cwd, bytes) = decode::read_str_from_slice(bytes).map_err(error_report)?; + let (session, bytes) = decode::read_str_from_slice(bytes).map_err(error_report)?; + let (hostname, bytes) = decode::read_str_from_slice(bytes).map_err(error_report)?; + + let mut bytes = Bytes::new(bytes); + + let (deleted_at, bytes) = match decode::read_u64(&mut bytes) { + Ok(unix) => (Some(unix), bytes.remaining_slice()), + // we accept null here + Err(ValueReadError::TypeMismatch(Marker::Null)) => (None, bytes.remaining_slice()), + Err(err) => return Err(error_report(err)), + }; + let (author, bytes) = Self::read_optional_string(bytes)?; + let (intent, bytes) = if nfields > 10 { + Self::read_optional_string(bytes)? + } else { + (None, bytes) + }; + + if !bytes.is_empty() { + bail!("trailing bytes in encoded history. malformed") + } + + Ok(Self { + id: id.to_owned().into(), + timestamp: OffsetDateTime::from_unix_timestamp_nanos(i128::from(timestamp))?, + duration, + exit, + command: command.to_owned(), + cwd: cwd.to_owned(), + session: session.to_owned(), + hostname: hostname.to_owned(), + author: author.unwrap_or_else(|| Self::author_from_hostname(hostname)), + intent, + deleted_at: deleted_at + .map(|t| OffsetDateTime::from_unix_timestamp_nanos(i128::from(t))) + .transpose()?, + }) + } + + pub(crate) fn deserialize(bytes: &[u8], version: &str) -> Result<Self> { + match version { + HISTORY_VERSION_V0 => Self::deserialize_v0(bytes), + HISTORY_VERSION_V1 => Self::deserialize_v1(bytes), + + _ => bail!("unknown version {version:?}"), + } + } + + /// Builder for a history entry that is captured via hook. + /// + /// This builder is used only at the `start` step of the hook, + /// so it doesn't have any fields which are known only after + /// the command is finished, such as `exit` or `duration`. + /// + /// ## Examples + /// ```rust + /// use crate::aclient::history::History; + /// + /// let history: History = History::capture() + /// .timestamp(time::OffsetDateTime::now_utc()) + /// .command("ls -la") + /// .cwd("/home/user") + /// .build() + /// .into(); + /// ``` + /// + /// Command without any required info cannot be captured, which is forced at compile time: + /// + /// ```compile_fail + /// use crate::aclient::history::History; + /// + /// // this will not compile because `cwd` is missing + /// let history: History = History::capture() + /// .timestamp(time::OffsetDateTime::now_utc()) + /// .command("ls -la") + /// .build() + /// .into(); + /// ``` + pub(crate) fn capture() -> builder::HistoryCapturedBuilder { + builder::HistoryCaptured::builder() + } + + /// Builder for a history entry that is captured via hook, and sent to the daemon. + /// + /// This builder is used only at the `start` step of the hook, + /// so it doesn't have any fields which are known only after + /// the command is finished, such as `exit` or `duration`. + /// + /// It does, however, include information that can usually be inferred. + /// + /// This is because the daemon we are sending a request to lacks the context of the command + /// + /// ## Examples + /// ```rust + /// use crate::aclient::history::History; + /// + /// let history: History = History::daemon() + /// .timestamp(time::OffsetDateTime::now_utc()) + /// .command("ls -la") + /// .cwd("/home/user") + /// .session("018deb6e8287781f9973ef40e0fde76b") + /// .hostname("computer:ellie") + /// .build() + /// .into(); + /// ``` + /// + /// Command without any required info cannot be captured, which is forced at compile time: + /// + /// ```compile_fail + /// use crate::aclient::history::History; + /// + /// // this will not compile because `hostname` is missing + /// let history: History = History::daemon() + /// .timestamp(time::OffsetDateTime::now_utc()) + /// .command("ls -la") + /// .cwd("/home/user") + /// .session("018deb6e8287781f9973ef40e0fde76b") + /// .build() + /// .into(); + /// ``` + pub(crate) fn daemon() -> builder::HistoryDaemonCaptureBuilder { + builder::HistoryDaemonCapture::builder() + } + + /// Builder for a history entry that is imported from the database. + /// + /// All fields are required, as they are all present in the database. + /// + /// ```compile_fail + /// use crate::aclient::history::History; + /// + /// // this will not compile because `id` field is missing + /// let history: History = History::from_db() + /// .timestamp(time::OffsetDateTime::now_utc()) + /// .command("ls -la".to_string()) + /// .cwd("/home/user".to_string()) + /// .exit(0) + /// .duration(100) + /// .session("somesession".to_string()) + /// .hostname("localhost".to_string()) + /// .author("user".to_string()) + /// .intent(None) + /// .deleted_at(None) + /// .build() + /// .into(); + /// ``` + pub(crate) fn from_db() -> builder::HistoryFromDbBuilder { + builder::HistoryFromDb::builder() + } + + pub(crate) fn success(&self) -> bool { + self.exit == 0 || self.duration == -1 + } + + pub(crate) fn should_save(&self, settings: &Settings) -> bool { + !(self.command.is_empty() + || settings.history_filter.is_match(&self.command) + || settings.cwd_filter.is_match(&self.cwd) + || (settings.secrets_filter && SECRET_PATTERNS_RE.is_match(&self.command))) + } +} + +#[cfg(test)] +mod tests { + use regex::RegexSet; + use time::macros::datetime; + + use crate::aclient::{history::HISTORY_VERSION, settings::Settings}; + + use super::History; + + // Test that we don't save history where necessary + #[test] + fn privacy_test() { + let settings = Settings { + cwd_filter: RegexSet::new(["^/supasecret"]).unwrap(), + history_filter: RegexSet::new(["^psql"]).unwrap(), + ..Settings::default() + }; + + let normal_command: History = History::capture() + .timestamp(time::OffsetDateTime::now_utc()) + .command("echo foo") + .cwd("/") + .build() + .into(); + + let with_space: History = History::capture() + .timestamp(time::OffsetDateTime::now_utc()) + .command(" echo bar") + .cwd("/") + .build() + .into(); + + let empty: History = History::capture() + .timestamp(time::OffsetDateTime::now_utc()) + .command("") + .cwd("/") + .build() + .into(); + + let stripe_key: History = History::capture() + .timestamp(time::OffsetDateTime::now_utc()) + .command("curl foo.com/bar?key=sk_test_1234567890abcdefghijklmnop") + .cwd("/") + .build() + .into(); + + let secret_dir: History = History::capture() + .timestamp(time::OffsetDateTime::now_utc()) + .command("echo ohno") + .cwd("/supasecret") + .build() + .into(); + + let with_psql: History = History::capture() + .timestamp(time::OffsetDateTime::now_utc()) + .command("psql") + .cwd("/supasecret") + .build() + .into(); + + assert!(normal_command.should_save(&settings)); + assert!(!with_space.should_save(&settings)); + assert!(!empty.should_save(&settings)); + assert!(!stripe_key.should_save(&settings)); + assert!(!secret_dir.should_save(&settings)); + assert!(!with_psql.should_save(&settings)); + } + + #[test] + fn disable_secrets() { + let settings = Settings { + secrets_filter: false, + ..Settings::new().unwrap() + }; + + let stripe_key: History = History::capture() + .timestamp(time::OffsetDateTime::now_utc()) + .command("curl foo.com/bar?key=sk_test_1234567890abcdefghijklmnop") + .cwd("/") + .build() + .into(); + + assert!(stripe_key.should_save(&settings)); + } + + #[test] + fn test_serialize_deserialize() { + let history = History { + id: "66d16cbee7cd47538e5c5b8b44e9006e".to_owned().into(), + timestamp: datetime!(2023-05-28 18:35:40.633872 +00:00), + duration: 49_206_000, + exit: 0, + command: "git status".to_owned(), + cwd: "/Users/conrad.ludgate/Documents/code/atuin".to_owned(), + session: "b97d9a306f274473a203d2eba41f9457".to_owned(), + hostname: "fvfg936c0kpf:conrad.ludgate".to_owned(), + author: "conrad.ludgate".to_owned(), + intent: None, + deleted_at: None, + }; + + let serialized = history.serialize().expect("failed to serialize history"); + assert_eq!( + &serialized.0[0..3], + [205, 0, 1], + "should encode as history v1" + ); + + let deserialized = History::deserialize(&serialized.0, HISTORY_VERSION) + .expect("failed to deserialize history"); + assert_eq!(history, deserialized); + } + + #[test] + fn test_serialize_deserialize_deleted() { + let history = History { + id: "66d16cbee7cd47538e5c5b8b44e9006e".to_owned().into(), + timestamp: datetime!(2023-05-28 18:35:40.633872 +00:00), + duration: 49_206_000, + exit: 0, + command: "git status".to_owned(), + cwd: "/Users/conrad.ludgate/Documents/code/atuin".to_owned(), + session: "b97d9a306f274473a203d2eba41f9457".to_owned(), + hostname: "fvfg936c0kpf:conrad.ludgate".to_owned(), + author: "conrad.ludgate".to_owned(), + intent: None, + deleted_at: Some(datetime!(2023-11-19 20:18 +00:00)), + }; + + let serialized = history.serialize().expect("failed to serialize history"); + + let deserialized = History::deserialize(&serialized.0, HISTORY_VERSION) + .expect("failed to deserialize history"); + + assert_eq!(history, deserialized); + } + + #[test] + fn test_serialize_deserialize_with_author_and_intent() { + let history = History { + id: "66d16cbee7cd47538e5c5b8b44e9006e".to_owned().into(), + timestamp: datetime!(2023-05-28 18:35:40.633872 +00:00), + duration: 49_206_000, + exit: 0, + command: "git status".to_owned(), + cwd: "/Users/conrad.ludgate/Documents/code/atuin".to_owned(), + session: "b97d9a306f274473a203d2eba41f9457".to_owned(), + hostname: "fvfg936c0kpf:conrad.ludgate".to_owned(), + author: "claude".to_owned(), + intent: Some("check repository status".to_owned()), + deleted_at: None, + }; + + let serialized = history.serialize().expect("failed to serialize history"); + let deserialized = History::deserialize(&serialized.0, HISTORY_VERSION) + .expect("failed to deserialize history"); + + assert_eq!(history, deserialized); + } + + #[test] + fn test_serialize_deserialize_version() { + // v0 + let bytes_v0 = [ + 205, 0, 0, 153, 217, 32, 54, 54, 100, 49, 54, 99, 98, 101, 101, 55, 99, 100, 52, 55, + 53, 51, 56, 101, 53, 99, 53, 98, 56, 98, 52, 52, 101, 57, 48, 48, 54, 101, 207, 23, 99, + 98, 117, 24, 210, 246, 128, 206, 2, 238, 210, 240, 0, 170, 103, 105, 116, 32, 115, 116, + 97, 116, 117, 115, 217, 42, 47, 85, 115, 101, 114, 115, 47, 99, 111, 110, 114, 97, 100, + 46, 108, 117, 100, 103, 97, 116, 101, 47, 68, 111, 99, 117, 109, 101, 110, 116, 115, + 47, 99, 111, 100, 101, 47, 97, 116, 117, 105, 110, 217, 32, 98, 57, 55, 100, 57, 97, + 51, 48, 54, 102, 50, 55, 52, 52, 55, 51, 97, 50, 48, 51, 100, 50, 101, 98, 97, 52, 49, + 102, 57, 52, 53, 55, 187, 102, 118, 102, 103, 57, 51, 54, 99, 48, 107, 112, 102, 58, + 99, 111, 110, 114, 97, 100, 46, 108, 117, 100, 103, 97, 116, 101, 192, + ]; + + let deserialized = History::deserialize(&bytes_v0, "v0"); + assert!(deserialized.is_ok()); + + let deserialized = History::deserialize(&bytes_v0, HISTORY_VERSION); + assert!(deserialized.is_err()); + + let current = History { + id: "66d16cbee7cd47538e5c5b8b44e9006e".to_owned().into(), + timestamp: datetime!(2023-05-28 18:35:40.633872 +00:00), + duration: 49_206_000, + exit: 0, + command: "git status".to_owned(), + cwd: "/Users/conrad.ludgate/Documents/code/atuin".to_owned(), + session: "b97d9a306f274473a203d2eba41f9457".to_owned(), + hostname: "fvfg936c0kpf:conrad.ludgate".to_owned(), + author: "conrad.ludgate".to_owned(), + intent: None, + deleted_at: None, + }; + + let bytes_v1 = current.serialize().expect("failed to serialize history"); + let deserialized = History::deserialize(&bytes_v1.0, HISTORY_VERSION); + assert!(deserialized.is_ok()); + + let deserialized = History::deserialize(&bytes_v1.0, "v0"); + assert!(deserialized.is_err()); + } +} diff --git a/crates/daemon/src/aclient/history/store.rs b/crates/daemon/src/aclient/history/store.rs new file mode 100644 index 00000000..db692590 --- /dev/null +++ b/crates/daemon/src/aclient/history/store.rs @@ -0,0 +1,437 @@ +use std::{collections::HashSet, fmt::Write, time::Duration}; + +use eyre::{Result, bail, eyre}; +use indicatif::{ProgressBar, ProgressState, ProgressStyle}; +use rmp::decode::Bytes; +use tracing::debug; + +use crate::aclient::{ + database::{ClientSqlite, current_context}, + record::{encryption::PASETO_V4, sqlite_store::SqliteStore}, +}; +use turtle_common::record::{DecryptedData, Host, HostId, Record, RecordId, RecordIdx}; + +use super::{HISTORY_TAG, HISTORY_VERSION, HISTORY_VERSION_V0, History, HistoryId}; + +#[derive(Debug, Clone)] +pub(crate) struct HistoryStore { + pub(crate) store: SqliteStore, + pub(crate) host_id: HostId, + pub(crate) encryption_key: [u8; 32], +} + +#[derive(Debug, Eq, PartialEq, Clone)] +pub(crate) enum HistoryRecord { + Create(History), // Create a history record + Delete(HistoryId), // Delete a history record, identified by ID +} + +impl HistoryRecord { + /// Serialize a history record, returning `DecryptedData` + /// The record will be of a certain type + /// We map those like so: + /// + /// `HistoryRecord::Create` -> 0 + /// `HistoryRecord::Delete`-> 1 + /// + /// This numeric identifier is then written as the first byte to the buffer. For history, we + /// append the serialized history right afterwards, to avoid having to handle serialization + /// twice. + /// + /// Deletion simply refers to the history by ID + pub(crate) fn serialize(&self) -> Result<DecryptedData> { + // probably don't actually need to use rmp here, but if we ever need to extend it, it's a + // nice wrapper around raw byte stuff + use rmp::encode; + + let mut output = vec![]; + + match self { + Self::Create(history) => { + // 0 -> a history create + encode::write_u8(&mut output, 0)?; + + let bytes = history.serialize()?; + + encode::write_bin(&mut output, &bytes.0)?; + } + Self::Delete(id) => { + // 1 -> a history delete + encode::write_u8(&mut output, 1)?; + encode::write_str(&mut output, id.0.as_str())?; + } + } + + Ok(DecryptedData(output)) + } + + pub(crate) fn deserialize(bytes: &DecryptedData, version: &str) -> Result<Self> { + use rmp::decode; + + fn error_report<E: std::fmt::Debug>(err: E) -> eyre::Report { + eyre!("{err:?}") + } + + let mut bytes = Bytes::new(&bytes.0); + + let record_type = decode::read_u8(&mut bytes).map_err(error_report)?; + + match record_type { + // 0 -> HistoryRecord::Create + 0 => { + // not super useful to us atm, but perhaps in the future + // written by write_bin above + decode::read_bin_len(&mut bytes).map_err(error_report)?; + + let record = History::deserialize(bytes.remaining_slice(), version)?; + + Ok(Self::Create(record)) + } + + // 1 -> HistoryRecord::Delete + 1 => { + let bytes = bytes.remaining_slice(); + let (id, bytes) = decode::read_str_from_slice(bytes).map_err(error_report)?; + + if !bytes.is_empty() { + bail!( + "trailing bytes decoding HistoryRecord::Delete - malformed? got {bytes:?}" + ); + } + + Ok(Self::Delete(id.to_string().into())) + } + + n => { + bail!("unknown HistoryRecord type {n}") + } + } + } +} + +impl HistoryStore { + pub(crate) fn new(store: SqliteStore, host_id: HostId, encryption_key: [u8; 32]) -> Self { + Self { + store, + host_id, + encryption_key, + } + } + + async fn push_record(&self, record: HistoryRecord) -> Result<(RecordId, RecordIdx)> { + let bytes = record.serialize()?; + let idx = self + .store + .last(self.host_id, HISTORY_TAG) + .await? + .map_or(0, |p| p.idx + 1); + + let record = Record::builder() + .host(Host::new(self.host_id)) + .version(HISTORY_VERSION.to_string()) + .tag(HISTORY_TAG.to_string()) + .idx(idx) + .data(bytes) + .build(); + + let id = record.id; + + self.store + .push(&record.encrypt::<PASETO_V4>(&self.encryption_key)) + .await?; + + Ok((id, idx)) + } + + async fn push_batch(&self, records: impl Iterator<Item = HistoryRecord>) -> Result<()> { + let mut ret = Vec::new(); + + let idx = self + .store + .last(self.host_id, HISTORY_TAG) + .await? + .map_or(0, |p| p.idx + 1); + + // Could probably _also_ do this as an iterator, but let's see how this is for now. + // optimizing for minimal sqlite transactions, this code can be optimised later + for (n, record) in records.enumerate() { + let bytes = record.serialize()?; + + let record = Record::builder() + .host(Host::new(self.host_id)) + .version(HISTORY_VERSION.to_string()) + .tag(HISTORY_TAG.to_string()) + .idx(idx + n as u64) + .data(bytes) + .build(); + + let record = record.encrypt::<PASETO_V4>(&self.encryption_key); + + ret.push(record); + } + + self.store.push_batch(ret.iter()).await?; + + Ok(()) + } + + pub(crate) async fn delete(&self, id: HistoryId) -> Result<(RecordId, RecordIdx)> { + let record = HistoryRecord::Delete(id); + + self.push_record(record).await + } + + /// Delete a batch of history entries via the record store. + /// Returns the record IDs so the caller can run `incremental_build` when ready. + pub(crate) async fn delete_entries( + &self, + entries: impl IntoIterator<Item = History>, + ) -> Result<Vec<RecordId>> { + let mut record_ids = Vec::new(); + for entry in entries { + let (id, _) = self.delete(entry.id).await?; + record_ids.push(id); + } + Ok(record_ids) + } + + pub(crate) async fn push(&self, history: History) -> Result<(RecordId, RecordIdx)> { + // TODO(ellie): move the history store to its own file + // it's tiny rn so fine as is + let record = HistoryRecord::Create(history); + + self.push_record(record).await + } + + pub(crate) async fn history(&self) -> Result<Vec<HistoryRecord>> { + // Atm this loads all history into memory + // Not ideal as that is potentially quite a lot, although history will be small. + let records = self.store.all_tagged(HISTORY_TAG).await?; + let mut ret = Vec::with_capacity(records.len()); + + for record in records { + let hist = match record.version.as_str() { + HISTORY_VERSION_V0 | HISTORY_VERSION => { + let version = record.version.clone(); + let decrypted = record.decrypt::<PASETO_V4>(&self.encryption_key)?; + + HistoryRecord::deserialize(&decrypted.data, version.as_str()) + } + version => bail!("unknown history version {version:?}"), + }?; + + ret.push(hist); + } + + Ok(ret) + } + + pub(crate) async fn build(&self, database: &ClientSqlite) -> Result<()> { + // I'd like to change how we rebuild and not couple this with the database, but need to + // consider the structure more deeply. This will be easy to change. + + // TODO(ellie): page or iterate this + let history = self.history().await?; + + // In theory we could flatten this here + // The current issue is that the database may have history in it already, from the old sync + // This didn't actually delete old history + // If we're sure we have a DB only maintained by the new store, we can flatten + // create/delete before we even get to sqlite + let mut creates = Vec::new(); + let mut deletes = Vec::new(); + + for i in history { + match i { + HistoryRecord::Create(h) => { + creates.push(h); + } + HistoryRecord::Delete(id) => { + deletes.push(id); + } + } + } + + database.save_bulk(&creates).await?; + database.delete_rows(&deletes).await?; + + Ok(()) + } + + pub(crate) async fn incremental_build( + &self, + database: &ClientSqlite, + ids: &[RecordId], + ) -> Result<()> { + for id in ids { + let record = self.store.get(*id).await; + + if let Ok(record) = record { + if record.tag != HISTORY_TAG { + continue; + } + + let version = record.version.clone(); + let decrypted = record.decrypt::<PASETO_V4>(&self.encryption_key)?; + let record = match version.as_str() { + HISTORY_VERSION_V0 | HISTORY_VERSION => { + HistoryRecord::deserialize(&decrypted.data, version.as_str())? + } + version => bail!("unknown history version {version:?}"), + }; + + match record { + HistoryRecord::Create(h) => { + // TODO: benchmark CPU time/memory tradeoff of batch commit vs one at a time + database.save(&h).await?; + } + HistoryRecord::Delete(id) => { + database.delete_rows(&[id]).await?; + } + } + } + } + + Ok(()) + } + + /// Get a list of history IDs that exist in the store + /// Note: This currently involves loading all history into memory. This is not going to be a + /// large amount in absolute terms, but do not all it in a hot loop. + pub(crate) async fn history_ids(&self) -> Result<HashSet<HistoryId>> { + let history = self.history().await?; + + let ret = history + .iter() + .map(|h| match h { + HistoryRecord::Create(h) => h.id.clone(), + HistoryRecord::Delete(id) => id.clone(), + }) + .collect::<HashSet<_>>(); + + Ok(ret) + } + + pub(crate) async fn init_store(&self, db: &ClientSqlite) -> Result<()> { + let pb = ProgressBar::new_spinner(); + pb.set_style( + ProgressStyle::with_template("{spinner:.blue} {msg}") + .unwrap() + .with_key("eta", |state: &ProgressState, w: &mut dyn Write| { + write!(w, "{:.1}s", state.eta().as_secs_f64()).unwrap(); + }) + .progress_chars("#>-"), + ); + pb.enable_steady_tick(Duration::from_millis(500)); + + pb.set_message("Fetching history from old database"); + + let context = current_context().await?; + let history = db.list(&[], &context, None, false, true).await?; + + pb.set_message("Fetching history already in store"); + let store_ids = self.history_ids().await?; + + pb.set_message("Converting old history to new store"); + let mut records = Vec::new(); + + for i in history { + debug!("loaded {}", i.id); + + if store_ids.contains(&i.id) { + debug!("skipping {} - already exists", i.id); + continue; + } + + if i.deleted_at.is_some() { + records.push(HistoryRecord::Delete(i.id)); + } else { + records.push(HistoryRecord::Create(i)); + } + } + + pb.set_message("Writing to db"); + + if !records.is_empty() { + self.push_batch(records.into_iter()).await?; + } + + pb.finish_with_message("Import complete"); + + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use turtle_common::record::DecryptedData; + use time::macros::datetime; + + use crate::aclient::history::{HISTORY_VERSION, store::HistoryRecord}; + + use super::History; + + #[test] + fn test_serialize_deserialize_create() { + let bytes = [ + 204, 0, 196, 147, 205, 0, 1, 154, 217, 32, 48, 49, 56, 99, 100, 52, 102, 101, 56, 49, + 55, 53, 55, 99, 100, 50, 97, 101, 101, 54, 53, 99, 100, 55, 56, 54, 49, 102, 57, 99, + 56, 49, 207, 23, 166, 251, 212, 181, 82, 0, 0, 100, 0, 162, 108, 115, 217, 41, 47, 85, + 115, 101, 114, 115, 47, 101, 108, 108, 105, 101, 47, 115, 114, 99, 47, 103, 105, 116, + 104, 117, 98, 46, 99, 111, 109, 47, 97, 116, 117, 105, 110, 115, 104, 47, 97, 116, 117, + 105, 110, 217, 32, 48, 49, 56, 99, 100, 52, 102, 101, 97, 100, 56, 57, 55, 53, 57, 55, + 56, 53, 50, 53, 50, 55, 97, 51, 49, 99, 57, 57, 56, 48, 53, 57, 170, 98, 111, 111, 112, + 58, 101, 108, 108, 105, 101, 192, 165, 101, 108, 108, 105, 101, + ]; + + let history = History { + id: "018cd4fe81757cd2aee65cd7861f9c81".to_owned().into(), + timestamp: datetime!(2024-01-04 00:00:00.000000 +00:00), + duration: 100, + exit: 0, + command: "ls".to_owned(), + cwd: "/Users/ellie/src/github.com/atuinsh/atuin".to_owned(), + session: "018cd4fead897597852527a31c998059".to_owned(), + hostname: "boop:ellie".to_owned(), + author: "ellie".to_owned(), + intent: None, + deleted_at: None, + }; + + let record = HistoryRecord::Create(history); + + let serialized = record.serialize().expect("failed to serialize history"); + assert_eq!(serialized.0, bytes); + + let deserialized = HistoryRecord::deserialize(&serialized, HISTORY_VERSION) + .expect("failed to deserialize HistoryRecord"); + assert_eq!(deserialized, record); + + // check the snapshot too + let deserialized = + HistoryRecord::deserialize(&DecryptedData(Vec::from(bytes)), HISTORY_VERSION) + .expect("failed to deserialize HistoryRecord"); + assert_eq!(deserialized, record); + } + + #[test] + fn test_serialize_deserialize_delete() { + let bytes = [ + 204, 1, 217, 32, 48, 49, 56, 99, 100, 52, 102, 101, 56, 49, 55, 53, 55, 99, 100, 50, + 97, 101, 101, 54, 53, 99, 100, 55, 56, 54, 49, 102, 57, 99, 56, 49, + ]; + let record = HistoryRecord::Delete("018cd4fe81757cd2aee65cd7861f9c81".to_string().into()); + + let serialized = record.serialize().expect("failed to serialize history"); + assert_eq!(serialized.0, bytes); + + let deserialized = HistoryRecord::deserialize(&serialized, HISTORY_VERSION) + .expect("failed to deserialize HistoryRecord"); + assert_eq!(deserialized, record); + + let deserialized = + HistoryRecord::deserialize(&DecryptedData(Vec::from(bytes)), HISTORY_VERSION) + .expect("failed to deserialize HistoryRecord"); + assert_eq!(deserialized, record); + } +} diff --git a/crates/daemon/src/aclient/meta.rs b/crates/daemon/src/aclient/meta.rs new file mode 100644 index 00000000..00dabcce --- /dev/null +++ b/crates/daemon/src/aclient/meta.rs @@ -0,0 +1,182 @@ +use std::path::Path; +use std::str::FromStr; +use std::time::Duration; + +use turtle_common::record::HostId; +use eyre::{Result, eyre}; +use sqlx::sqlite::{SqliteConnectOptions, SqliteJournalMode, SqlitePool, SqlitePoolOptions}; +use time::{OffsetDateTime, format_description::well_known::Rfc3339}; +use tokio::sync::OnceCell; +use tracing::debug; +use uuid::Uuid; + +const KEY_HOST_ID: &str = "host_id"; +const KEY_LAST_SYNC: &str = "last_sync_time"; + +pub(crate) struct MetaStore { + pool: SqlitePool, + cached_host_id: OnceCell<HostId>, +} + +impl MetaStore { + pub(crate) async fn new(path: impl AsRef<Path>, timeout: f64) -> Result<Self> { + let path = path.as_ref(); + let path_str = path + .as_os_str() + .to_str() + .ok_or_else(|| eyre!("meta database path is not valid UTF-8: {path:?}"))?; + debug!("opening meta sqlite database at {path:?}"); + + let is_memory = path_str.contains(":memory:"); + + if !is_memory + && !path.exists() + && let Some(dir) = path.parent() + { + fs_err::create_dir_all(dir)?; + } + + // Use DELETE journal mode instead of WAL. This is a small, infrequently- + // written KV store — WAL's concurrency benefits aren't needed, and DELETE + // mode avoids creating auxiliary -wal/-shm files that complicate + // permission handling. + let opts = SqliteConnectOptions::from_str(path_str)? + .journal_mode(SqliteJournalMode::Delete) + .optimize_on_close(true, None) + .create_if_missing(true); + + let pool = SqlitePoolOptions::new() + .acquire_timeout(Duration::from_secs_f64(timeout)) + .connect_with(opts) + .await?; + + sqlx::migrate!("./db/client-meta-migrations") + .run(&pool) + .await?; + + // Session tokens are stored in this database, so restrict permissions. + #[cfg(unix)] + if !is_memory { + use std::os::unix::fs::PermissionsExt; + std::fs::set_permissions(path, std::fs::Permissions::from_mode(0o600))?; + } + + let store = Self { + pool, + cached_host_id: OnceCell::const_new(), + }; + + Ok(store) + } + + // Generic key-value operations + + pub(crate) async fn get(&self, key: &str) -> Result<Option<String>> { + let row: Option<(String,)> = sqlx::query_as("SELECT value FROM meta WHERE key = ?1") + .bind(key) + .fetch_optional(&self.pool) + .await?; + + Ok(row.map(|r| r.0)) + } + + pub(crate) async fn set(&self, key: &str, value: &str) -> Result<()> { + sqlx::query( + " + INSERT INTO meta (key, value, updated_at) + VALUES (?1, ?2, strftime('%s', 'now')) + ON CONFLICT(key) DO UPDATE + SET value = ?2, updated_at = strftime('%s', 'now') + ", + ) + .bind(key) + .bind(value) + .execute(&self.pool) + .await?; + + Ok(()) + } + + // Typed accessors + + pub(crate) async fn host_id(&self) -> Result<HostId> { + self.cached_host_id + .get_or_try_init(|| async { + if let Some(id) = self.get(KEY_HOST_ID).await? { + let parsed = Uuid::from_str(id.as_str()) + .map_err(|e| eyre!("failed to parse host ID: {e}"))?; + return Ok(HostId(parsed)); + } + + let uuid = turtle_common::utils::uuid_v7(); + self.set(KEY_HOST_ID, uuid.as_simple().to_string().as_ref()) + .await?; + + Ok(HostId(uuid)) + }) + .await + .copied() + } + + pub(crate) async fn last_sync(&self) -> Result<OffsetDateTime> { + match self.get(KEY_LAST_SYNC).await? { + Some(v) => Ok(OffsetDateTime::parse(v.as_str(), &Rfc3339)?), + None => Ok(OffsetDateTime::UNIX_EPOCH), + } + } + + pub(crate) async fn save_sync_time(&self) -> Result<()> { + self.set( + KEY_LAST_SYNC, + OffsetDateTime::now_utc().format(&Rfc3339)?.as_str(), + ) + .await + } +} + +#[cfg(test)] +mod tests { + use super::{MetaStore, OffsetDateTime}; + + async fn new_test_store() -> MetaStore { + MetaStore::new("sqlite::memory:", 2.0).await.unwrap() + } + + #[tokio::test] + async fn test_get_set_delete() { + let store = new_test_store().await; + + assert_eq!(store.get("foo").await.unwrap(), None); + + store.set("foo", "bar").await.unwrap(); + assert_eq!(store.get("foo").await.unwrap(), Some("bar".to_string())); + + store.set("foo", "baz").await.unwrap(); + assert_eq!(store.get("foo").await.unwrap(), Some("baz".to_string())); + + store.delete("foo").await.unwrap(); + assert_eq!(store.get("foo").await.unwrap(), None); + } + + #[tokio::test] + async fn test_host_id_generation_and_stability() { + let store = new_test_store().await; + + let id1 = store.host_id().await.unwrap(); + let id2 = store.host_id().await.unwrap(); + + assert_eq!(id1, id2, "host_id should be stable across calls"); + } + + #[tokio::test] + async fn test_sync_time() { + let store = new_test_store().await; + + let t = store.last_sync().await.unwrap(); + assert_eq!(t, OffsetDateTime::UNIX_EPOCH); + + store.save_sync_time().await.unwrap(); + let t = store.last_sync().await.unwrap(); + assert!(t > OffsetDateTime::UNIX_EPOCH); + } +} diff --git a/crates/daemon/src/aclient/mod.rs b/crates/daemon/src/aclient/mod.rs new file mode 100644 index 00000000..3f3709a9 --- /dev/null +++ b/crates/daemon/src/aclient/mod.rs @@ -0,0 +1,11 @@ +pub mod database; +pub mod history; +pub mod record; +pub mod settings; + +pub(crate) mod encryption; +pub(crate) mod meta; +pub(crate) mod utils; +pub(crate) mod ordering; +pub(crate) mod secrets; +pub(crate) mod api_client; diff --git a/crates/daemon/src/aclient/ordering.rs b/crates/daemon/src/aclient/ordering.rs new file mode 100644 index 00000000..84001f52 --- /dev/null +++ b/crates/daemon/src/aclient/ordering.rs @@ -0,0 +1,31 @@ +use minspan::minspan; + +use super::{history::History, settings::SearchMode}; + +pub(crate) fn reorder_fuzzy(mode: SearchMode, query: &str, res: Vec<History>) -> Vec<History> { + match mode { + SearchMode::Fuzzy => reorder(query, |x| &x.command, res), + _ => res, + } +} + +#[expect(clippy::needless_pass_by_value, reason = "makes things easier")] +fn reorder<F, A>(query: &str, f: F, res: Vec<A>) -> Vec<A> +where + F: Fn(&A) -> &String, + A: Clone, +{ + let mut r = res.clone(); + let qvec = &query.chars().collect(); + r.sort_by_cached_key(|h| { + // TODO for fzf search we should sum up scores for each matched term + // + // The fallback is a little unfortunate: when we are asked to match a query that is found nowhere, + // we don't want to return a None, as the comparison behaviour would put the worst matches + // at the front. Therefore, we'll return a set of indices that are one larger than the longest + // possible legitimate match. This is meaningless except as a comparison. + let (from, to) = minspan::span(qvec, &(f(h).chars().collect())).unwrap_or((0, res.len())); + 1 + to - from + }); + r +} diff --git a/crates/daemon/src/aclient/record/encryption.rs b/crates/daemon/src/aclient/record/encryption.rs new file mode 100644 index 00000000..6851a99f --- /dev/null +++ b/crates/daemon/src/aclient/record/encryption.rs @@ -0,0 +1,379 @@ +use turtle_common::record::{ + AdditionalData, DecryptedData, EncryptedData, Encryption, HostId, RecordId, RecordIdx, +}; +use base64::{Engine, engine::general_purpose}; +use eyre::{Context, Result, ensure}; +use rusty_paserk::{Key, KeyId, Local, PieWrappedKey}; +use rusty_paseto::core::{ + ImplicitAssertion, Key as DataKey, Local as LocalPurpose, Paseto, PasetoNonce, Payload, V4, +}; +use serde::{Deserialize, Serialize}; + +/// Use PASETO V4 Local encryption using the additional data as an implicit assertion. +#[expect(non_camel_case_types)] +pub(crate) struct PASETO_V4; + +/* +Why do we use a random content-encryption key? +Originally I was planning on using a derived key for encryption based on additional data. +This would be a lot more secure than using the master key directly. + +However, there's an established norm of using a random key. This scheme might be otherwise known as +- client-side encryption +- envelope encryption +- key wrapping + +A HSM (Hardware Security Module) provider, eg: AWS, Azure, GCP, or even a physical device like a YubiKey +will have some keys that they keep to themselves. These keys never leave their physical hardware. +If they never leave the hardware, then encrypting large amounts of data means giving them the data and waiting. +This is not a practical solution. Instead, generate a unique key for your data, encrypt that using your HSM +and then store that with your data. + +See + - <https://docs.aws.amazon.com/wellarchitected/latest/financial-services-industry-lens/use-envelope-encryption-with-customer-master-keys.html> + - <https://cloud.google.com/kms/docs/envelope-encryption> + - <https://learn.microsoft.com/en-us/azure/storage/blobs/client-side-encryption?tabs=dotnet#encryption-and-decryption-via-the-envelope-technique> + - <https://www.yubico.com/gb/product/yubihsm-2-fips/> + - <https://cheatsheetseries.owasp.org/cheatsheets/Cryptographic_Storage_Cheat_Sheet.html#encrypting-stored-keys> + +Why would we care? In the past we have received some requests for company solutions. If in future we can configure a +KMS service with little effort, then that would solve a lot of issues for their security team. + +Even for personal use, if a user is not comfortable with sharing keys between hosts, +GCP HSM costs $1/month and $0.03 per 10,000 key operations. Assuming an active user runs +1000 atuin records a day, that would only cost them $1 and 10 cent a month. + +Additionally, key rotations are much simpler using this scheme. Rotating a key is as simple as re-encrypting the CEK, and not the message contents. +This makes it very fast to rotate a key in bulk. + +For future reference, with asymmetric encryption, you can encrypt the CEK without the HSM's involvement, but decrypting +will need the HSM. This allows the encryption path to still be extremely fast (no network calls) but downloads/decryption +that happens in the background can make the network calls to the HSM +*/ + +impl Encryption for PASETO_V4 { + fn re_encrypt( + mut data: EncryptedData, + _ad: AdditionalData<'_>, + old_key: &[u8; 32], + new_key: &[u8; 32], + ) -> Result<EncryptedData> { + let cek = Self::decrypt_cek(&data.content_encryption_key, old_key)?; + data.content_encryption_key = Self::encrypt_cek(cek, new_key); + Ok(data) + } + + fn encrypt(data: DecryptedData, ad: AdditionalData<'_>, key: &[u8; 32]) -> EncryptedData { + // generate a random key for this entry + // aka content-encryption-key (CEK) + let random_key = Key::<V4, Local>::new_os_random(); + + let assertions = Assertions::from(ad).encode(); + + // build the payload and encrypt the token + let payload = serde_json::to_string(&AtuinPayload { + data: general_purpose::URL_SAFE_NO_PAD.encode(data.0), + }) + .expect("json encoding can't fail"); + let nonce = DataKey::<32>::try_new_random().expect("could not source from random"); + let nonce = PasetoNonce::<V4, LocalPurpose>::from(&nonce); + + let token = Paseto::<V4, LocalPurpose>::builder() + .set_payload(Payload::from(payload.as_str())) + .set_implicit_assertion(ImplicitAssertion::from(assertions.as_str())) + .try_encrypt(&random_key.into(), &nonce) + .expect("error encrypting atuin data"); + + EncryptedData { + data: token, + content_encryption_key: Self::encrypt_cek(random_key, key), + } + } + + fn decrypt( + data: EncryptedData, + ad: AdditionalData<'_>, + key: &[u8; 32], + ) -> Result<DecryptedData> { + let token = data.data; + let cek = Self::decrypt_cek(&data.content_encryption_key, key)?; + + // encode the implicit assertions + let assertions = Assertions::from(ad).encode(); + + // decrypt the payload with the footer and implicit assertions + let payload = Paseto::<V4, LocalPurpose>::try_decrypt( + &token, + &cek.into(), + None, + ImplicitAssertion::from(&*assertions), + ) + .context("could not decrypt entry")?; + + let payload: AtuinPayload = serde_json::from_str(&payload)?; + let data = general_purpose::URL_SAFE_NO_PAD.decode(payload.data)?; + Ok(DecryptedData(data)) + } +} + +impl PASETO_V4 { + fn decrypt_cek(wrapped_cek: &str, key: &[u8; 32]) -> Result<Key<V4, Local>> { + let wrapping_key = Key::<V4, Local>::from_bytes(*key); + + // let wrapping_key = PasetoSymmetricKey::from(Key::from(key)); + + let AtuinFooter { kid, wpk } = serde_json::from_str(wrapped_cek) + .context("wrapped cek did not contain the correct contents")?; + + // check that the wrapping key matches the required key to decrypt. + // In future, we could support multiple keys and use this key to + // look up the key rather than only allow one key. + // For now though we will only support the one key and key rotation will + // have to be a hard reset + let current_kid = wrapping_key.to_id(); + + ensure!( + current_kid == kid, + "attempting to decrypt with incorrect key. currently using {current_kid}, expecting {kid}" + ); + + // decrypt the random key + Ok(wpk.unwrap_key(&wrapping_key)?) + } + + fn encrypt_cek(cek: Key<V4, Local>, key: &[u8; 32]) -> String { + // aka key-encryption-key (KEK) + let wrapping_key = Key::<V4, Local>::from_bytes(*key); + let kid = wrapping_key.to_id(); + + // wrap the random key so we can decrypt it later + let wrapped_cek = AtuinFooter { + wpk: cek.wrap_pie(&wrapping_key), + kid, + }; + serde_json::to_string(&wrapped_cek).expect("could not serialize wrapped cek") + } +} + +#[derive(Serialize, Deserialize)] +struct AtuinPayload { + data: String, +} + +#[derive(Serialize, Deserialize)] +/// Well-known footer claims for decrypting. This is not encrypted but is stored in the record. +/// <https://github.com/paseto-standard/paseto-spec/blob/master/docs/02-Implementation-Guide/04-Claims.md#optional-footer-claims> +struct AtuinFooter { + /// Wrapped key + wpk: PieWrappedKey<V4, Local>, + /// ID of the key which was used to wrap + kid: KeyId<V4, Local>, +} + +/// Used in the implicit assertions. This is not encrypted and not stored in the data blob. +// This cannot be changed, otherwise it breaks the authenticated encryption. +#[derive(Debug, Copy, Clone, Serialize)] +struct Assertions<'a> { + id: &'a RecordId, + idx: &'a RecordIdx, + version: &'a str, + tag: &'a str, + host: &'a HostId, +} + +impl<'a> From<AdditionalData<'a>> for Assertions<'a> { + fn from(ad: AdditionalData<'a>) -> Self { + Self { + id: ad.id, + version: ad.version, + tag: ad.tag, + host: ad.host, + idx: ad.idx, + } + } +} + +impl Assertions<'_> { + fn encode(&self) -> String { + serde_json::to_string(self).expect("could not serialize implicit assertions") + } +} + +#[cfg(test)] +mod tests { + use turtle_common::{ + record::{Host, Record}, + utils::uuid_v7, + }; + + use super::{ + AdditionalData, DecryptedData, Encryption, HostId, Key, Local, PASETO_V4, RecordId, V4, + }; + + #[test] + fn round_trip() { + let key = Key::<V4, Local>::new_os_random(); + + let ad = AdditionalData { + id: &RecordId(uuid_v7()), + version: "v0", + tag: "kv", + host: &HostId(uuid_v7()), + idx: &0, + }; + + let data = DecryptedData(vec![1, 2, 3, 4]); + + let encrypted = PASETO_V4::encrypt(data.clone(), ad, &key.to_bytes()); + let decrypted = PASETO_V4::decrypt(encrypted, ad, &key.to_bytes()).unwrap(); + assert_eq!(decrypted, data); + } + + #[test] + fn same_entry_different_output() { + let key = Key::<V4, Local>::new_os_random(); + + let ad = AdditionalData { + id: &RecordId(uuid_v7()), + version: "v0", + tag: "kv", + host: &HostId(uuid_v7()), + idx: &0, + }; + + let data = DecryptedData(vec![1, 2, 3, 4]); + + let encrypted = PASETO_V4::encrypt(data.clone(), ad, &key.to_bytes()); + let encrypted2 = PASETO_V4::encrypt(data, ad, &key.to_bytes()); + + assert_ne!( + encrypted.data, encrypted2.data, + "re-encrypting the same contents should have different output due to key randomization" + ); + } + + #[test] + fn cannot_decrypt_different_key() { + let key = Key::<V4, Local>::new_os_random(); + let fake_key = Key::<V4, Local>::new_os_random(); + + let ad = AdditionalData { + id: &RecordId(uuid_v7()), + version: "v0", + tag: "kv", + host: &HostId(uuid_v7()), + idx: &0, + }; + + let data = DecryptedData(vec![1, 2, 3, 4]); + + let encrypted = PASETO_V4::encrypt(data, ad, &key.to_bytes()); + drop(PASETO_V4::decrypt(encrypted, ad, &fake_key.to_bytes()).unwrap_err()); + } + + #[test] + fn cannot_decrypt_different_id() { + let key = Key::<V4, Local>::new_os_random(); + + let ad = AdditionalData { + id: &RecordId(uuid_v7()), + version: "v0", + tag: "kv", + host: &HostId(uuid_v7()), + idx: &0, + }; + + let data = DecryptedData(vec![1, 2, 3, 4]); + + let encrypted = PASETO_V4::encrypt(data, ad, &key.to_bytes()); + + let ad = AdditionalData { + id: &RecordId(uuid_v7()), + ..ad + }; + drop(PASETO_V4::decrypt(encrypted, ad, &key.to_bytes()).unwrap_err()); + } + + #[test] + fn re_encrypt_round_trip() { + let key1 = Key::<V4, Local>::new_os_random(); + let key2 = Key::<V4, Local>::new_os_random(); + + let ad = AdditionalData { + id: &RecordId(uuid_v7()), + version: "v0", + tag: "kv", + host: &HostId(uuid_v7()), + idx: &0, + }; + + let data = DecryptedData(vec![1, 2, 3, 4]); + + let encrypted1 = PASETO_V4::encrypt(data.clone(), ad, &key1.to_bytes()); + let encrypted2 = + PASETO_V4::re_encrypt(encrypted1.clone(), ad, &key1.to_bytes(), &key2.to_bytes()) + .unwrap(); + + // we only re-encrypt the content keys + assert_eq!(encrypted1.data, encrypted2.data); + assert_ne!( + encrypted1.content_encryption_key, + encrypted2.content_encryption_key + ); + + let decrypted = PASETO_V4::decrypt(encrypted2, ad, &key2.to_bytes()).unwrap(); + + assert_eq!(decrypted, data); + } + + #[test] + fn full_record_round_trip() { + let key = [0x55; 32]; + let record = Record::builder() + .id(RecordId(uuid_v7())) + .version("v0".to_owned()) + .tag("kv".to_owned()) + .host(Host::new(HostId(uuid_v7()))) + .timestamp(1_687_244_806_000_000) + .data(DecryptedData(vec![1, 2, 3, 4])) + .idx(0) + .build(); + + let encrypted = record.encrypt::<PASETO_V4>(&key); + + assert!(!encrypted.data.data.is_empty()); + assert!(!encrypted.data.content_encryption_key.is_empty()); + + let decrypted = encrypted.decrypt::<PASETO_V4>(&key).unwrap(); + + assert_eq!(decrypted.data.0, [1, 2, 3, 4]); + } + + #[test] + fn full_record_round_trip_fail() { + let key = [0x55; 32]; + let record = Record::builder() + .id(RecordId(uuid_v7())) + .version("v0".to_owned()) + .tag("kv".to_owned()) + .host(Host::new(HostId(uuid_v7()))) + .timestamp(1_687_244_806_000_000) + .data(DecryptedData(vec![1, 2, 3, 4])) + .idx(0) + .build(); + + let encrypted = record.encrypt::<PASETO_V4>(&key); + + let mut enc1 = encrypted.clone(); + enc1.host = Host::new(HostId(uuid_v7())); + let _ = enc1 + .decrypt::<PASETO_V4>(&key) + .expect_err("tampering with the host should result in auth failure"); + + let mut enc2 = encrypted; + enc2.id = RecordId(uuid_v7()); + let _ = enc2 + .decrypt::<PASETO_V4>(&key) + .expect_err("tampering with the id should result in auth failure"); + } +} diff --git a/crates/daemon/src/aclient/record/mod.rs b/crates/daemon/src/aclient/record/mod.rs new file mode 100644 index 00000000..2ace26f5 --- /dev/null +++ b/crates/daemon/src/aclient/record/mod.rs @@ -0,0 +1,4 @@ +pub(crate) mod encryption; +pub(crate) mod sync; + +pub mod sqlite_store; diff --git a/crates/daemon/src/aclient/record/sqlite_store.rs b/crates/daemon/src/aclient/record/sqlite_store.rs new file mode 100644 index 00000000..f2fc9d84 --- /dev/null +++ b/crates/daemon/src/aclient/record/sqlite_store.rs @@ -0,0 +1,563 @@ +// Here we are using sqlite as a pretty dumb store, and will not be running any complex queries. +// Multiple stores of multiple types are all stored in one chonky table (for now), and we just index +// by tag/host + +use std::path::Path; +use std::str::FromStr; + +use eyre::{Result, eyre}; +use fs_err as fs; + +use sqlx::{ + Row, + sqlite::{SqliteConnectOptions, SqliteJournalMode, SqlitePool, SqliteRow}, +}; +use tracing::debug; + +use crate::aclient::utils::setup_db; +use turtle_common::record::{ + EncryptedData, Host, HostId, Record, RecordId, RecordIdx, RecordStatus, +}; +use turtle_common::utils; +use uuid::Uuid; + +use super::encryption::PASETO_V4; + +#[derive(Debug, Clone)] +pub struct SqliteStore { + pool: SqlitePool, +} + +impl SqliteStore { + pub async fn new(path: impl AsRef<Path>, timeout: f64) -> Result<Self> { + fn mk_opts(path: &str) -> sqlx::Result<SqliteConnectOptions> { + let opts = SqliteConnectOptions::from_str(path)? + .journal_mode(SqliteJournalMode::Wal) + .foreign_keys(true) + .create_if_missing(true); + + Ok(opts) + } + + let path = path.as_ref(); + + debug!("opening sqlite database at {path:?}"); + + if utils::broken_symlink(path) { + eprintln!( + "Atuin: Sqlite db path ({}) is a broken symlink. Unable to read or create replacement.", + path.display() + ); + std::process::exit(1); + } + + if !path.exists() + && let Some(dir) = path.parent() + { + fs::create_dir_all(dir)?; + } + + let pool = setup_db!(path, timeout, mk_opts, "./db/client-record-migrations").await?; + + Ok(Self { pool }) + } + + async fn save_raw( + tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>, + r: &Record<EncryptedData>, + ) -> Result<()> { + // In sqlite, we are "limited" to i64. But that is still fine, until 2262. + sqlx::query( + "insert or ignore into store(id, idx, host, tag, timestamp, version, data, cek) + values(?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8)", + ) + .bind(r.id.0.as_hyphenated().to_string()) + .bind(r.idx as i64) + .bind(r.host.id.0.as_hyphenated().to_string()) + .bind(r.tag.as_str()) + .bind(r.timestamp as i64) + .bind(r.version.as_str()) + .bind(r.data.data.as_str()) + .bind(r.data.content_encryption_key.as_str()) + .execute(&mut **tx) + .await?; + + Ok(()) + } + + #[expect( + clippy::needless_pass_by_value, + reason = "this is used in a place with fixed function signature" + )] + fn query_row(row: SqliteRow) -> Record<EncryptedData> { + let idx: i64 = row.get("idx"); + let timestamp: i64 = row.get("timestamp"); + + // tbh at this point things are pretty fucked so just panic + let id = Uuid::from_str(row.get("id")).expect("invalid id UUID format in sqlite DB"); + let host = Uuid::from_str(row.get("host")).expect("invalid host UUID format in sqlite DB"); + + Record { + id: RecordId(id), + idx: idx as u64, + host: Host::new(HostId(host)), + timestamp: timestamp as u64, + tag: row.get("tag"), + version: row.get("version"), + data: EncryptedData { + data: row.get("data"), + content_encryption_key: row.get("cek"), + }, + } + } + + 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) + } +} + +/// A record store stores records +/// In more detail - we tend to need to process this into _another_ format to actually query it. +/// As is, the record store is intended as the source of truth for arbitrary data, which could +/// be shell history, kvs, etc. +impl SqliteStore { + /// Push a record + pub(crate) async fn push(&self, record: &Record<EncryptedData>) -> Result<()> { + self.push_batch(std::iter::once(record)).await + } + + /// Push a batch of records, all in one transaction + pub(crate) async fn push_batch( + &self, + records: impl Iterator<Item = &Record<EncryptedData>> + Send + Sync, + ) -> Result<()> { + let mut tx = self.pool.begin().await?; + + for record in records { + Self::save_raw(&mut tx, record).await?; + } + + tx.commit().await?; + + Ok(()) + } + + pub(crate) async fn get(&self, id: RecordId) -> Result<Record<EncryptedData>> { + let res = sqlx::query("select * from store where store.id = ?1") + .bind(id.0.as_hyphenated().to_string()) + .map(Self::query_row) + .fetch_one(&self.pool) + .await?; + + Ok(res) + } + + pub(crate) async fn delete(&self, id: RecordId) -> Result<()> { + sqlx::query("delete from store where id = ?1") + .bind(id.0.as_hyphenated().to_string()) + .execute(&self.pool) + .await?; + + Ok(()) + } + + pub(crate) async fn delete_all(&self) -> Result<()> { + sqlx::query("delete from store").execute(&self.pool).await?; + + Ok(()) + } + + pub(crate) async fn last( + &self, + host: HostId, + tag: &str, + ) -> Result<Option<Record<EncryptedData>>> { + let res = + sqlx::query("select * from store where host=?1 and tag=?2 order by idx desc limit 1") + .bind(host.0.as_hyphenated().to_string()) + .bind(tag) + .map(Self::query_row) + .fetch_one(&self.pool) + .await; + + match res { + Err(sqlx::Error::RowNotFound) => Ok(None), + Err(e) => Err(eyre!("an error occurred: {}", e)), + Ok(record) => Ok(Some(record)), + } + } + + pub(crate) async fn first( + &self, + host: HostId, + tag: &str, + ) -> Result<Option<Record<EncryptedData>>> { + self.idx(host, tag, 0).await + } + + pub(crate) async fn len_tag(&self, tag: &str) -> Result<u64> { + let res: Result<(i64,), sqlx::Error> = + sqlx::query_as("select count(*) from store where tag=?1") + .bind(tag) + .fetch_one(&self.pool) + .await; + match res { + Err(e) => Err(eyre!("failed to fetch local store len: {}", e)), + Ok(v) => Ok(v.0 as u64), + } + } + + /// Get the next `limit` records, after and including the given index + pub(crate) async fn next( + &self, + host: HostId, + tag: &str, + idx: RecordIdx, + limit: u64, + ) -> Result<Vec<Record<EncryptedData>>> { + let res = sqlx::query( + "select * from store where idx >= ?1 and host = ?2 and tag = ?3 order by idx asc limit ?4", + ) + .bind(idx as i64) + .bind(host.0.as_hyphenated().to_string()) + .bind(tag) + .bind(limit as i64) + .map(Self::query_row) + .fetch_all(&self.pool) + .await?; + + Ok(res) + } + + /// Get the first record for a given host and tag + pub(crate) async fn idx( + &self, + host: HostId, + tag: &str, + idx: RecordIdx, + ) -> Result<Option<Record<EncryptedData>>> { + let res = sqlx::query("select * from store where idx = ?1 and host = ?2 and tag = ?3") + .bind(idx as i64) + .bind(host.0.as_hyphenated().to_string()) + .bind(tag) + .map(Self::query_row) + .fetch_one(&self.pool) + .await; + + match res { + Err(sqlx::Error::RowNotFound) => Ok(None), + Err(e) => Err(eyre!("an error occurred: {}", e)), + Ok(v) => Ok(Some(v)), + } + } + + pub(crate) async fn status(&self) -> Result<RecordStatus> { + let mut status = RecordStatus::new(); + + let res: Result<Vec<(String, String, i64)>, sqlx::Error> = + sqlx::query_as("select host, tag, max(idx) from store group by host, tag") + .fetch_all(&self.pool) + .await; + + let res = match res { + Err(e) => return Err(eyre!("failed to fetch local store status: {}", e)), + Ok(v) => v, + }; + + for i in res { + let host = HostId( + Uuid::from_str(i.0.as_str()).expect("failed to parse uuid for local store status"), + ); + + status.set_raw(host, i.1, i.2 as u64); + } + + Ok(status) + } + + /// Get all records for a given tag + pub(crate) async fn all_tagged(&self, tag: &str) -> Result<Vec<Record<EncryptedData>>> { + let res = sqlx::query("select * from store where tag = ?1 order by timestamp asc") + .bind(tag) + .map(Self::query_row) + .fetch_all(&self.pool) + .await?; + + Ok(res) + } + + /// Reencrypt every single item in this store with a new key + /// Be careful - this may mess with sync. + pub(crate) 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(()) + } + + /// Verify that every record in this store can be decrypted with the current key + /// Someday maybe also check each tag/record can be deserialized, but not for now. + pub(crate) async fn verify(&self, key: &[u8; 32]) -> Result<()> { + let all = self.load_all().await?; + + all.into_iter() + .map(|record| record.decrypt::<PASETO_V4>(key)) + .collect::<Result<Vec<_>>>()?; + + Ok(()) + } + + /// Verify that every record in this store can be decrypted with the current key + /// Someday maybe also check each tag/record can be deserialized, but not for now. + pub(crate) async fn purge(&self, key: &[u8; 32]) -> Result<()> { + let all = self.load_all().await?; + + for record in &all { + if record.clone().decrypt::<PASETO_V4>(key).is_ok() { + continue; + } + + println!( + "Failed to decrypt {}, deleting", + record.id.0.as_hyphenated() + ); + + self.delete(record.id).await?; + } + + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use crate::{ + atuin_client::{ + encryption::generate_encoded_key, record::encryption::PASETO_V4, + settings::test_local_timeout, + }, + atuin_common::{ + record::{DecryptedData, EncryptedData, Host, HostId, Record}, + utils::uuid_v7, + }, + }; + + use super::SqliteStore; + + fn test_record() -> Record<EncryptedData> { + Record::builder() + .host(Host::new(HostId(uuid_v7()))) + .version("v1".into()) + .tag(uuid_v7().simple().to_string()) + .data(EncryptedData { + data: "1234".into(), + content_encryption_key: "1234".into(), + }) + .idx(0) + .build() + } + + #[tokio::test] + async fn create_db() { + let db = SqliteStore::new(":memory:", test_local_timeout()).await; + + assert!( + db.is_ok(), + "db could not be created, {:?}", + db.err().unwrap() + ); + } + + #[tokio::test] + async fn push_record() { + let db = SqliteStore::new(":memory:", test_local_timeout()) + .await + .unwrap(); + let record = test_record(); + + db.push(&record).await.expect("failed to insert record"); + } + + #[tokio::test] + async fn get_record() { + let db = SqliteStore::new(":memory:", test_local_timeout()) + .await + .unwrap(); + let record = test_record(); + db.push(&record).await.unwrap(); + + let new_record = db.get(record.id).await.expect("failed to fetch record"); + + assert_eq!(record, new_record, "records are not equal"); + } + + #[tokio::test] + async fn last() { + let db = SqliteStore::new(":memory:", test_local_timeout()) + .await + .unwrap(); + let record = test_record(); + db.push(&record).await.unwrap(); + + let last = db + .last(record.host.id, record.tag.as_str()) + .await + .expect("failed to get store len"); + + assert_eq!( + last.unwrap().id, + record.id, + "expected to get back the same record that was inserted" + ); + } + + #[tokio::test] + async fn first() { + let db = SqliteStore::new(":memory:", test_local_timeout()) + .await + .unwrap(); + let record = test_record(); + db.push(&record).await.unwrap(); + + let first = db + .first(record.host.id, record.tag.as_str()) + .await + .expect("failed to get store len"); + + assert_eq!( + first.unwrap().id, + record.id, + "expected to get back the same record that was inserted" + ); + } + + #[tokio::test] + async fn len() { + let db = SqliteStore::new(":memory:", test_local_timeout()) + .await + .unwrap(); + let record = test_record(); + db.push(&record).await.unwrap(); + + let len = db + .len(record.host.id, record.tag.as_str()) + .await + .expect("failed to get store len"); + + assert_eq!(len, 1, "expected length of 1 after insert"); + } + + #[tokio::test] + async fn len_tag() { + let db = SqliteStore::new(":memory:", test_local_timeout()) + .await + .unwrap(); + let record = test_record(); + db.push(&record).await.unwrap(); + + let len = db + .len_tag(record.tag.as_str()) + .await + .expect("failed to get store len"); + + assert_eq!(len, 1, "expected length of 1 after insert"); + } + + #[tokio::test] + async fn re_encrypt() { + let store = SqliteStore::new(":memory:", test_local_timeout()) + .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 with 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); + } +} diff --git a/crates/daemon/src/aclient/record/sync.rs b/crates/daemon/src/aclient/record/sync.rs new file mode 100644 index 00000000..94764f67 --- /dev/null +++ b/crates/daemon/src/aclient/record/sync.rs @@ -0,0 +1,456 @@ +// do a sync :O +use std::{cmp::Ordering, fmt::Write}; + +use eyre::{OptionExt, Result}; +use thiserror::Error; +use tracing::error; + +use super::encryption::PASETO_V4; +use crate::aclient::record::sqlite_store::SqliteStore; +use crate::aclient::{api_client::Client, settings::Settings}; + +use turtle_common::record::{Diff, HostId, RecordId, RecordIdx, RecordStatus}; +use indicatif::{ProgressBar, ProgressState, ProgressStyle}; + +#[derive(Error, Debug)] +pub(crate) enum SyncError { + #[error("an issue with the local database occurred: {msg:?}")] + LocalStoreError { msg: String }, + + #[error("something has gone wrong with the sync logic: {msg:?}")] + SyncLogicError { msg: String }, + + #[error("operational error: {msg:?}")] + OperationalError { msg: String }, + + #[error("a request to the sync server failed: {msg}")] + RemoteRequestError { msg: String }, + + #[error( + "the encryption key on this machine does not match the data on the server. \ + this usually means a new machine was set up without copying the existing key. \ + to fix: run `atuin key` on a machine that already syncs correctly, then run \ + `atuin store rekey <key>` on this machine with the value from the other machine" + )] + WrongKey, +} + +#[derive(Debug, Eq, PartialEq)] +pub(crate) enum Operation { + // Either upload or download until the states matches the below + Upload { + local: RecordIdx, + remote: Option<RecordIdx>, + host: HostId, + tag: String, + }, + Download { + local: Option<RecordIdx>, + remote: RecordIdx, + host: HostId, + tag: String, + }, + Noop { + host: HostId, + tag: String, + }, +} + +pub(crate) fn build_client(settings: &Settings) -> Result<Client<'_>, SyncError> { + Client::new( + &settings.sync.address, + settings.network_connect_timeout, + settings.network_timeout, + settings + .sync + .user_id() + .map_err(|e| SyncError::RemoteRequestError { msg: e.to_string() })? + .ok_or_eyre("No sync user-id set") + .map_err(|e| SyncError::RemoteRequestError { msg: e.to_string() })?, + ) + .map_err(|e| SyncError::OperationalError { msg: e.to_string() }) +} + +pub(crate) async fn diff( + client: &Client<'_>, + store: &SqliteStore, +) -> Result<(Vec<Diff>, RecordStatus), SyncError> { + let local_index = store + .status() + .await + .map_err(|e| SyncError::LocalStoreError { msg: e.to_string() })?; + + let remote_index = client + .record_status() + .await + .map_err(|e| SyncError::RemoteRequestError { msg: e.to_string() })?; + + let diff = local_index.diff(&remote_index); + + Ok((diff, remote_index)) +} + +// Take a diff, along with a local store, and resolve it into a set of operations. +// With the store as context, we can determine if a tail exists locally or not and therefore if it needs uploading or download. +// In theory this could be done as a part of the diffing stage, but it's easier to reason +// about and test this way +pub(crate) fn operations( + diffs: Vec<Diff>, + _store: &SqliteStore, +) -> Result<Vec<Operation>, SyncError> { + let mut operations = Vec::with_capacity(diffs.len()); + + for diff in diffs { + let op = match (diff.local, diff.remote) { + // We both have it! Could be either. Compare. + (Some(local), Some(remote)) => match local.cmp(&remote) { + Ordering::Equal => Operation::Noop { + host: diff.host, + tag: diff.tag, + }, + Ordering::Greater => Operation::Upload { + local, + remote: Some(remote), + host: diff.host, + tag: diff.tag, + }, + Ordering::Less => Operation::Download { + local: Some(local), + remote, + host: diff.host, + tag: diff.tag, + }, + }, + + // Remote has it, we don't. Gotta be download + (None, Some(remote)) => Operation::Download { + local: None, + remote, + host: diff.host, + tag: diff.tag, + }, + + // We have it, remote doesn't. Gotta be upload. + (Some(local), None) => Operation::Upload { + local, + remote: None, + host: diff.host, + tag: diff.tag, + }, + + // something is pretty fucked. + (None, None) => { + return Err(SyncError::SyncLogicError { + msg: String::from( + "diff has nothing for local or remote - (host, tag) does not exist", + ), + }); + } + }; + + operations.push(op); + } + + // sort them - purely so we have a stable testing order, and can rely on + // same input = same output + // We can sort by ID so long as we continue to use UUIDv7 or something + // with the same properties + + operations.sort_by_key(|op| match op { + Operation::Noop { host, tag } => (0, *host, tag.clone()), + + Operation::Upload { host, tag, .. } => (1, *host, tag.clone()), + + Operation::Download { host, tag, .. } => (2, *host, tag.clone()), + }); + + Ok(operations) +} + +async fn sync_upload( + store: &SqliteStore, + client: &Client<'_>, + host: HostId, + tag: String, + local: RecordIdx, + remote: Option<RecordIdx>, + page_size: u64, +) -> Result<i64, SyncError> { + let remote = remote.unwrap_or(0); + let expected = local - remote; + let mut progress = 0; + + let pb = ProgressBar::new(expected); + pb.set_style(ProgressStyle::with_template("{spinner:.green} [{elapsed_precise}] [{wide_bar:.cyan/blue}] {human_pos}/{human_len} ({eta})") + .unwrap() + .with_key("eta", |state: &ProgressState, w: &mut dyn Write| write!(w, "{:.1}s", state.eta().as_secs_f64()).unwrap()) + .progress_chars("#>-")); + + println!( + "Uploading {} records to {}/{}", + expected, + host.0.as_simple(), + tag + ); + + loop { + let page = store + .next(host, tag.as_str(), remote + progress, page_size) + .await + .map_err(|e| { + error!("failed to read upload page: {e:?}"); + + SyncError::LocalStoreError { msg: e.to_string() } + })?; + + if page.is_empty() { + break; + } + + client.post_records(&page).await.map_err(|e| { + error!("failed to post records: {e:?}"); + + SyncError::RemoteRequestError { msg: e.to_string() } + })?; + + progress += page.len() as u64; + pb.set_position(progress); + + if progress >= expected { + break; + } + } + + pb.finish_with_message("Uploaded records"); + + Ok(progress as i64) +} + +async fn sync_download( + store: &SqliteStore, + client: &Client<'_>, + host: HostId, + tag: String, + local: Option<RecordIdx>, + remote: RecordIdx, + page_size: u64, +) -> Result<Vec<RecordId>, SyncError> { + let local = local.unwrap_or(0); + let expected = remote - local; + let mut progress = 0; + let mut ret = Vec::new(); + + println!( + "Downloading {} records from {}/{}", + expected, + host.0.as_simple(), + tag + ); + + let pb = ProgressBar::new(expected); + pb.set_style(ProgressStyle::with_template("{spinner:.green} [{elapsed_precise}] [{wide_bar:.cyan/blue}] {human_pos}/{human_len} ({eta})") + .unwrap() + .with_key("eta", |state: &ProgressState, w: &mut dyn Write| write!(w, "{:.1}s", state.eta().as_secs_f64()).unwrap()) + .progress_chars("#>-")); + + loop { + let page = client + .next_records(host, tag.clone(), local + progress, page_size) + .await + .map_err(|e| SyncError::RemoteRequestError { msg: e.to_string() })?; + + if page.is_empty() { + break; + } + + store + .push_batch(page.iter()) + .await + .map_err(|e| SyncError::LocalStoreError { msg: e.to_string() })?; + + ret.extend(page.iter().map(|f| f.id)); + + progress += page.len() as u64; + pb.set_position(progress); + + if progress >= expected { + break; + } + } + + pb.finish_with_message("Downloaded records"); + + Ok(ret) +} + +pub(crate) async fn sync_remote( + client: &Client<'_>, + operations: Vec<Operation>, + local_store: &SqliteStore, + page_size: u64, +) -> Result<(i64, Vec<RecordId>), SyncError> { + let mut uploaded = 0; + let mut downloaded = Vec::new(); + + // this can totally run in parallel, but lets get it working first + for i in operations { + match i { + Operation::Upload { + host, + tag, + local, + remote, + } => { + uploaded += + sync_upload(local_store, client, host, tag, local, remote, page_size).await?; + } + + Operation::Download { + host, + tag, + local, + remote, + } => { + let mut d = + sync_download(local_store, client, host, tag, local, remote, page_size).await?; + downloaded.append(&mut d); + } + + Operation::Noop { .. } => (), + } + } + + Ok((uploaded, downloaded)) +} + +pub(crate) async fn check_encryption_key( + client: &Client<'_>, + remote_index: &RecordStatus, + encryption_key: &[u8; 32], +) -> Result<(), SyncError> { + let sample = remote_index + .hosts + .iter() + .flat_map(|(host, tags)| tags.keys().map(move |tag| (*host, tag.clone()))) + .next(); + + let Some((host, tag)) = sample else { + return Ok(()); + }; + + let records = client + .next_records(host, tag, 0, 1) + .await + .map_err(|e| SyncError::RemoteRequestError { msg: e.to_string() })?; + + let Some(record) = records.into_iter().next() else { + return Ok(()); + }; + + record.decrypt::<PASETO_V4>(encryption_key).map_err(|err| { + error!("Wrong key error: {err}"); + SyncError::WrongKey + })?; + + Ok(()) +} + +pub(crate) async fn sync( + settings: &Settings, + store: &SqliteStore, + encryption_key: &[u8; 32], +) -> Result<(i64, Vec<RecordId>), SyncError> { + let client = build_client(settings)?; + let (diff, remote_index) = diff(&client, store).await?; + + // Bail before mutating either side if the local key can't read the remote. + check_encryption_key(&client, &remote_index, encryption_key).await?; + + let operations = operations(diff, store)?; + let (uploaded, downloaded) = sync_remote(&client, operations, store, 100).await?; + + Ok((uploaded, downloaded)) +} + +#[cfg(test)] +mod tests { + use crate::aclient::record::sync::Operation; + use turtle_common::record::{Diff, EncryptedData, HostId, Record}; + + use crate::aclient::{ + record::{ + sqlite_store::SqliteStore, + sync::{self}, + }, + settings::test_local_timeout, + }; + + fn test_record() -> Record<EncryptedData> { + Record::builder() + .host(turtle_common::record::Host::new(HostId( + turtle_common::utils::uuid_v7(), + ))) + .version("v1".into()) + .tag(turtle_common::utils::uuid_v7().simple().to_string()) + .data(EncryptedData { + data: String::new(), + content_encryption_key: String::new(), + }) + .idx(0) + .build() + } + + // Take a list of local records, and a list of remote records. + // Return the local database, and a diff of local/remote, ready to build + // ops + async fn build_test_diff( + local_records: Vec<Record<EncryptedData>>, + remote_records: Vec<Record<EncryptedData>>, + ) -> (SqliteStore, Vec<Diff>) { + let local_store = SqliteStore::new(":memory:", test_local_timeout()) + .await + .expect("failed to open in memory sqlite"); + let remote_store = SqliteStore::new(":memory:", test_local_timeout()) + .await + .expect("failed to open in memory sqlite"); // "remote" + + for i in local_records { + local_store.push(&i).await.unwrap(); + } + + for i in remote_records { + remote_store.push(&i).await.unwrap(); + } + + let local_index = local_store.status().await.unwrap(); + let remote_index = remote_store.status().await.unwrap(); + + let diff = local_index.diff(&remote_index); + + (local_store, diff) + } + + #[tokio::test] + async fn test_basic_diff() { + // a diff where local is ahead of remote. nothing else. + + let record = test_record(); + let (store, diff) = build_test_diff(vec![record.clone()], vec![]).await; + + assert_eq!(diff.len(), 1); + + let operations = sync::operations(diff, &store).unwrap(); + + assert_eq!(operations.len(), 1); + + assert_eq!( + operations[0], + Operation::Upload { + host: record.host.id, + tag: record.tag, + local: record.idx, + remote: None, + } + ); + } +} diff --git a/crates/daemon/src/aclient/secrets.rs b/crates/daemon/src/aclient/secrets.rs new file mode 100644 index 00000000..08d24339 --- /dev/null +++ b/crates/daemon/src/aclient/secrets.rs @@ -0,0 +1,223 @@ +// This file will probably trigger a lot of scanners. Sorry. + +use regex::RegexSet; +use std::sync::LazyLock; + +#[cfg(test)] +pub(crate) enum TestValue<'a> { + Single(&'a str), + Multiple(&'a [&'a str]), +} + +#[cfg(test)] +type SpType<'a> = &'a [(&'a str, &'a str, TestValue<'a>)]; + +#[cfg(not(test))] +type SpType<'a> = &'a [(&'a str, &'a str)]; + +/// A list of `(name, regex, test)`, where `test` should match against `regex`. +pub(crate) static SECRET_PATTERNS: SpType<'_> = &[ + ( + "AWS Access Key ID", + "A[KS]IA[0-9A-Z]{16}", + #[cfg(test)] + TestValue::Single("AKIAIOSFODNN7EXAMPLE"), + ), + ( + "AWS Secret Access Key env var", + "AWS_SECRET_ACCESS_KEY", + #[cfg(test)] + TestValue::Single("AWS_SECRET_ACCESS_KEY=KEYDATA"), + ), + ( + "AWS Session Token env var", + "AWS_SESSION_TOKEN", + #[cfg(test)] + TestValue::Single("AWS_SESSION_TOKEN=KEYDATA"), + ), + ( + "Microsoft Azure secret access key env var", + "AZURE_.*_KEY", + #[cfg(test)] + TestValue::Single("export AZURE_STORAGE_ACCOUNT_KEY=KEYDATA"), + ), + ( + "Google cloud platform key env var", + "GOOGLE_SERVICE_ACCOUNT_KEY", + #[cfg(test)] + TestValue::Single("export GOOGLE_SERVICE_ACCOUNT_KEY=KEYDATA"), + ), + ( + "Atuin login", + r"atuin\s+login", + #[cfg(test)] + TestValue::Single( + "atuin login -u mycoolusername -p mycoolpassword -k \"lots of random words\"", + ), + ), + ( + "GitHub PAT (old)", + "ghp_[a-zA-Z0-9]{36}", + #[cfg(test)] + TestValue::Single("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"), // legit, I expired it + ), + ( + "GitHub PAT (new)", + "gh1_[A-Za-z0-9]{21}_[A-Za-z0-9]{59}|github_pat_[0-9][A-Za-z0-9]{21}_[A-Za-z0-9]{59}", + #[cfg(test)] + TestValue::Multiple(&[ + "gh1_1234567890abcdefghijk_1234567890abcdefghijklmnopqrstuvwxyz1234567890abcdefghijklm", + "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", // also legit, also expired + ]), + ), + ( + "GitHub OAuth Access Token", + "gho_[A-Za-z0-9]{36}", + #[cfg(test)] + TestValue::Single("gho_1234567890abcdefghijklmnopqrstuvwx000"), // not a real token + ), + ( + "GitHub OAuth Access Token (user)", + "ghu_[A-Za-z0-9]{36}", + #[cfg(test)] + TestValue::Single("ghu_1234567890abcdefghijklmnopqrstuvwx000"), // not a real token + ), + ( + "GitHub App Installation Access Token", + "ghs_[A-Za-z0-9._-]{36,}", + #[cfg(test)] + TestValue::Multiple(&[ + "ghs_1234567890abcdefghijklmnopqrstuvwx000", // not a real token + "ghs_abc-def.ghi_jklMNOP0123456789qrstuv-wxyzABCD", // new token format, fake data + ]), + ), + ( + "GitHub Refresh Token", + "ghr_[A-Za-z0-9]{76}", + #[cfg(test)] + TestValue::Single( + "ghr_1234567890abcdefghijklmnopqrstuvwx1234567890abcdefghijklmnopqrstuvwx1234567890abcdefghijklmnopqrstuvwx", + ), // not a real token + ), + ( + "GitHub App Installation Access Token v1", + "v1\\.[0-9A-Fa-f]{40}", + #[cfg(test)] + TestValue::Single("v1.1234567890abcdef1234567890abcdef12345678"), // not a real token + ), + ( + "GitLab PAT", + "glpat-[a-zA-Z0-9_]{20}", + #[cfg(test)] + TestValue::Single("glpat-RkE_BG5p_bbjML21WSfy"), + ), + ( + "Slack OAuth v2 bot", + "xoxb-[0-9]{11}-[0-9]{11}-[0-9a-zA-Z]{24}", + #[cfg(test)] + TestValue::Single("xoxb-17653672481-19874698323-pdFZKVeTuE8sk7oOcBrzbqgy"), + ), + ( + "Slack OAuth v2 user token", + "xoxp-[0-9]{11}-[0-9]{11}-[0-9a-zA-Z]{24}", + #[cfg(test)] + TestValue::Single("xoxp-17653672481-19874698323-pdFZKVeTuE8sk7oOcBrzbqgy"), + ), + ( + "Slack webhook", + "T[a-zA-Z0-9_]{8}/B[a-zA-Z0-9_]{8}/[a-zA-Z0-9_]{24}", + #[cfg(test)] + TestValue::Single( + "https://hooks.slack.com/services/T00000000/B00000000/XXXXXXXXXXXXXXXXXXXXXXXX", + ), + ), + ( + "Stripe test key", + "sk_test_[0-9a-zA-Z]{24}", + #[cfg(test)] + TestValue::Single("sk_test_1234567890abcdefghijklmnop"), + ), + ( + "Stripe live key", + "sk_live_[0-9a-zA-Z]{24}", + #[cfg(test)] + TestValue::Single("sk_live_1234567890abcdefghijklmnop"), + ), + ( + "Netlify authentication token", + "nf[pcoub]_[0-9a-zA-Z]{36}", + #[cfg(test)] + TestValue::Single("nfp_nBh7BdJxUwyaBBwFzpyD29MMFT6pZ9wq5634"), + ), + ( + "npm token", + "npm_[A-Za-z0-9]{36}", + #[cfg(test)] + TestValue::Single("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"), + ), + ( + "Pulumi personal access token", + "pul-[0-9a-f]{40}", + #[cfg(test)] + TestValue::Single("pul-683c2770662c51d960d72ec27613be7653c5cb26"), + ), +]; + +/// The `regex` expressions from [`SECRET_PATTERNS`] compiled into a `RegexSet`. +pub(crate) static SECRET_PATTERNS_RE: LazyLock<RegexSet> = LazyLock::new(|| { + let exprs = SECRET_PATTERNS.iter().map(|f| f.1); + RegexSet::new(exprs).expect("Failed to build secrets regex") +}); + +#[cfg(test)] +mod tests { + use regex::Regex; + + use crate::aclient::secrets::{SECRET_PATTERNS, TestValue}; + + #[test] + fn test_secrets() { + for (name, regex, test) in SECRET_PATTERNS { + let re = + Regex::new(regex).unwrap_or_else(|_| panic!("Failed to compile regex for {name}")); + + match test { + TestValue::Single(test) => { + assert!(re.is_match(test), "{name} test failed!"); + } + TestValue::Multiple(tests) => { + for test_str in tests.iter() { + assert!( + re.is_match(test_str), + "{name} test with value \"{test_str}\" failed!" + ); + } + } + } + } + } + + #[test] + fn test_secrets_embedded() { + for (name, regex, test) in SECRET_PATTERNS { + let re = + Regex::new(regex).unwrap_or_else(|_| panic!("Failed to compile regex for {name}")); + + match test { + TestValue::Single(test) => { + let embedded = format!("some random text {test} some more random text"); + assert!(re.is_match(&embedded), "{name} embedded test failed!"); + } + TestValue::Multiple(tests) => { + for test_str in tests.iter() { + let embedded = format!("some random text {test_str} some more random text"); + assert!( + re.is_match(&embedded), + "{name} embedded test with value \"{test_str}\" failed!" + ); + } + } + } + } + } +} diff --git a/crates/daemon/src/aclient/settings/meta.rs b/crates/daemon/src/aclient/settings/meta.rs new file mode 100644 index 00000000..7993ef6d --- /dev/null +++ b/crates/daemon/src/aclient/settings/meta.rs @@ -0,0 +1,17 @@ +use serde::{Deserialize, Serialize}; + +#[derive(Debug, Serialize, Deserialize, Clone)] +pub(crate) struct Settings { + pub(crate) db_path: String, +} + +impl Default for Settings { + fn default() -> Self { + let dir = turtle_common::utils::data_dir(); + let path = dir.join("meta.db"); + + Self { + db_path: path.to_string_lossy().to_string(), + } + } +} diff --git a/crates/daemon/src/aclient/settings/mod.rs b/crates/daemon/src/aclient/settings/mod.rs new file mode 100644 index 00000000..58afd17c --- /dev/null +++ b/crates/daemon/src/aclient/settings/mod.rs @@ -0,0 +1,1544 @@ +use crypto_secretbox::Key; +use std::{ + collections::HashMap, fmt, fs::read_to_string, path::PathBuf, str::FromStr, sync::OnceLock, +}; +use tokio::sync::OnceCell; +use tracing::info; +use uuid::Uuid; + +use crate::aclient::encryption::decode_key; +use clap::ValueEnum; +use config::{ + Config, ConfigBuilder, Environment, File as ConfigFile, FileFormat, builder::DefaultState, +}; +use eyre::{Context, Error, Result, bail, eyre}; +use fs_err::create_dir_all; +use regex::RegexSet; +use serde::{Deserialize, Serialize}; +use serde_with::DeserializeFromStr; +use time::{OffsetDateTime, UtcOffset, format_description::FormatItem, macros::format_description}; +use turtle_common::record::HostId; +use turtle_common::utils; + +static DATA_DIR: OnceLock<PathBuf> = OnceLock::new(); +static META_CONFIG: OnceLock<(String, f64)> = OnceLock::new(); +static META_STORE: OnceCell<crate::aclient::meta::MetaStore> = OnceCell::const_new(); + +pub(crate) mod meta; +pub(crate) mod watcher; + +#[derive(Clone, Debug, Deserialize, Copy, ValueEnum, PartialEq, Serialize)] +pub(crate) enum SearchMode { + #[serde(rename = "prefix")] + Prefix, + + #[serde(rename = "fulltext")] + #[clap(aliases = &["fulltext"])] + FullText, + + #[serde(rename = "fuzzy")] + Fuzzy, + + #[serde(rename = "skim")] + Skim, + + #[serde(rename = "daemon-fuzzy")] + #[clap(aliases = &["daemon-fuzzy"])] + DaemonFuzzy, +} + +impl SearchMode { + pub(crate) fn as_str(self) -> &'static str { + match self { + Self::Prefix => "PREFIX", + Self::FullText => "FULLTXT", + Self::Fuzzy => "FUZZY", + Self::Skim => "SKIM", + Self::DaemonFuzzy => "DAEMON", + } + } + pub(crate) fn next(self, settings: &Settings) -> Self { + match self { + Self::Prefix => Self::FullText, + // if the user is using skim, we go to skim + Self::FullText if settings.search_mode == Self::Skim => Self::Skim, + // if the user is using daemon-fuzzy, we go to daemon-fuzzy + Self::FullText if settings.search_mode == Self::DaemonFuzzy => Self::DaemonFuzzy, + // otherwise fuzzy. + Self::FullText => Self::Fuzzy, + Self::Fuzzy | Self::Skim | Self::DaemonFuzzy => Self::Prefix, + } + } +} + +#[derive(Clone, Debug, Deserialize, Copy, PartialEq, Eq, ValueEnum, Serialize)] +pub(crate) enum FilterMode { + #[serde(rename = "global")] + Global = 0, + + #[serde(rename = "host")] + Host = 1, + + #[serde(rename = "session")] + Session = 2, + + #[serde(rename = "directory")] + Directory = 3, + + #[serde(rename = "workspace")] + Workspace = 4, + + #[serde(rename = "session-preload")] + SessionPreload = 5, +} + +impl FilterMode { + pub(crate) fn as_str(self) -> &'static str { + match self { + Self::Global => "GLOBAL", + Self::Host => "HOST", + Self::Session => "SESSION", + Self::Directory => "DIRECTORY", + Self::Workspace => "WORKSPACE", + Self::SessionPreload => "SESSION+", + } + } +} + +#[derive(Clone, Debug, Deserialize, Copy, Serialize)] +pub(crate) enum ExitMode { + #[serde(rename = "return-original")] + ReturnOriginal, + + #[serde(rename = "return-query")] + ReturnQuery, +} + +// FIXME: Can use upstream Dialect enum if https://github.com/stevedonovan/chrono-english/pull/16 is merged +// FIXME: Above PR was merged, but dependency was changed to interim (fork of chrono-english) in the ... interim +#[derive(Clone, Debug, Deserialize, Copy, Serialize)] +pub(crate) enum Dialect { + #[serde(rename = "us")] + Us, + + #[serde(rename = "uk")] + Uk, +} + +impl From<Dialect> for interim::Dialect { + fn from(d: Dialect) -> Self { + match d { + Dialect::Uk => Self::Uk, + Dialect::Us => Self::Us, + } + } +} + +/// Type wrapper around `time::UtcOffset` to support a wider variety of timezone formats. +/// +/// Note that the parsing of this struct needs to be done before starting any +/// multithreaded runtime, otherwise it will fail on most Unix systems. +/// +/// See: <https://github.com/atuinsh/atuin/pull/1517#discussion_r1447516426> +#[derive(Clone, Copy, Debug, Eq, PartialEq, DeserializeFromStr, Serialize)] +pub(crate) struct Timezone(pub(crate) UtcOffset); +impl fmt::Display for Timezone { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + self.0.fmt(f) + } +} +/// format: <+|-><hour>[:<minute>[:<second>]] +static OFFSET_FMT: &[FormatItem<'_>] = format_description!( + "[offset_hour sign:mandatory padding:none][optional [:[offset_minute padding:none][optional [:[offset_second padding:none]]]]]" +); +impl FromStr for Timezone { + type Err = Error; + + fn from_str(s: &str) -> Result<Self> { + // local timezone + if matches!(s.to_lowercase().as_str(), "l" | "local") { + // There have been some timezone issues, related to errors fetching it on some + // platforms + // Rather than fail to start, fallback to UTC. The user should still be able to specify + // their timezone manually in the config file. + let offset = UtcOffset::current_local_offset().unwrap_or(UtcOffset::UTC); + return Ok(Self(offset)); + } + + if matches!(s.to_lowercase().as_str(), "0" | "utc") { + let offset = UtcOffset::UTC; + return Ok(Self(offset)); + } + + // offset from UTC + if let Ok(offset) = UtcOffset::parse(s, OFFSET_FMT) { + return Ok(Self(offset)); + } + + // IDEA: Currently named timezones are not supported, because the well-known crate + // for this is `chrono_tz`, which is not really interoperable with the datetime crate + // that we currently use - `time`. If ever we migrate to using `chrono`, this would + // be a good feature to add. + + bail!(r#""{s}" is not a valid timezone spec"#) + } +} + +#[derive(Clone, Debug, Deserialize, Copy, Serialize)] +pub(crate) enum Style { + #[serde(rename = "auto")] + Auto, + + #[serde(rename = "full")] + Full, + + #[serde(rename = "compact")] + Compact, +} + +#[derive(Clone, Debug, Deserialize, Copy, Serialize)] +pub(crate) enum WordJumpMode { + #[serde(rename = "emacs")] + Emacs, + + #[serde(rename = "subl")] + Subl, +} + +#[derive(Clone, Debug, Deserialize, Copy, PartialEq, Eq, ValueEnum, Serialize)] +pub(crate) enum KeymapMode { + #[serde(rename = "emacs")] + Emacs, + + #[serde(rename = "vim-normal")] + VimNormal, + + #[serde(rename = "vim-insert")] + VimInsert, + + #[serde(rename = "auto")] + Auto, +} + +// We want to translate the config to crossterm::cursor::SetCursorStyle, but +// the original type does not implement trait serde::Deserialize unfortunately. +// It seems impossible to implement Deserialize for external types when it is +// used in HashMap (https://stackoverflow.com/questions/67142663). We instead +// define an adapter type. +#[derive(Clone, Debug, Deserialize, Copy, PartialEq, Eq, ValueEnum, Serialize)] +pub(crate) enum CursorStyle { + #[serde(rename = "default")] + DefaultUserShape, + + #[serde(rename = "blink-block")] + BlinkingBlock, + + #[serde(rename = "steady-block")] + SteadyBlock, + + #[serde(rename = "blink-underline")] + BlinkingUnderScore, + + #[serde(rename = "steady-underline")] + SteadyUnderScore, + + #[serde(rename = "blink-bar")] + BlinkingBar, + + #[serde(rename = "steady-bar")] + SteadyBar, +} + +#[derive(Clone, Debug, Deserialize, Serialize)] +pub(crate) struct Stats { + #[serde(default = "Stats::common_prefix_default")] + pub(crate) common_prefix: Vec<String>, // sudo, etc. commands we want to strip off + #[serde(default = "Stats::common_subcommands_default")] + pub(crate) common_subcommands: Vec<String>, // kubectl, commands we should consider subcommands for + #[serde(default = "Stats::ignored_commands_default")] + pub(crate) ignored_commands: Vec<String>, // cd, ls, etc. commands we want to completely hide from stats +} + +impl Stats { + fn common_prefix_default() -> Vec<String> { + vec!["sudo", "doas"].into_iter().map(String::from).collect() + } + + fn common_subcommands_default() -> Vec<String> { + vec![ + "apt", + "cargo", + "composer", + "dnf", + "docker", + "dotnet", + "git", + "go", + "ip", + "jj", + "kubectl", + "nix", + "nmcli", + "npm", + "pecl", + "pnpm", + "podman", + "port", + "systemctl", + "tmux", + "yarn", + ] + .into_iter() + .map(String::from) + .collect() + } + + fn ignored_commands_default() -> Vec<String> { + vec![] + } +} + +impl Default for Stats { + fn default() -> Self { + Self { + common_prefix: Self::common_prefix_default(), + common_subcommands: Self::common_subcommands_default(), + ignored_commands: Self::ignored_commands_default(), + } + } +} + +#[derive(Clone, Debug, Deserialize, Default, Serialize)] +#[expect(clippy::struct_excessive_bools)] +pub(crate) struct Keys { + pub(crate) scroll_exits: bool, + pub(crate) exit_past_line_start: bool, + pub(crate) accept_past_line_end: bool, + pub(crate) accept_past_line_start: bool, + pub(crate) accept_with_backspace: bool, + pub(crate) prefix: String, +} + +impl Keys { + /// The standard default values for all `[keys]` options. + /// These match the config defaults set in `builder_with_data_dir()`. + pub(crate) fn standard_defaults() -> Self { + Self { + scroll_exits: true, + exit_past_line_start: true, + accept_past_line_end: true, + accept_past_line_start: false, + accept_with_backspace: false, + prefix: "a".to_string(), + } + } +} + +/// A single rule within a conditional keybinding config. +#[derive(Clone, Debug, Deserialize, Serialize)] +pub(crate) struct KeyRuleConfig { + /// Optional condition expression (e.g. "cursor-at-start", "input-empty && no-results"). + /// If absent, the rule always matches. + #[serde(default)] + pub(crate) when: Option<String>, + /// The action to perform (e.g. "exit", "cursor-left", "accept"). + pub(crate) action: String, +} + +/// A keybinding config value: either a simple action string or an ordered list of conditional rules. +#[derive(Clone, Debug, Deserialize, Serialize)] +#[serde(untagged)] +pub(crate) enum KeyBindingConfig { + /// Simple unconditional binding: `"ctrl-c" = "return-original"` + Simple(String), + /// Conditional binding: `"left" = [{ when = "cursor-at-start", action = "exit" }, { action = "cursor-left" }]` + Rules(Vec<KeyRuleConfig>), +} + +/// User-facing keymap configuration. Each mode maps key strings to bindings. +/// Keys present here override the defaults for that key; unmentioned keys keep defaults. +#[derive(Clone, Debug, Deserialize, Serialize, Default)] +pub(crate) struct KeymapConfig { + #[serde(default)] + pub(crate) emacs: HashMap<String, KeyBindingConfig>, + #[serde(default, rename = "vim-normal")] + pub(crate) vim_normal: HashMap<String, KeyBindingConfig>, + #[serde(default, rename = "vim-insert")] + pub(crate) vim_insert: HashMap<String, KeyBindingConfig>, + #[serde(default)] + pub(crate) inspector: HashMap<String, KeyBindingConfig>, + #[serde(default)] + pub(crate) prefix: HashMap<String, KeyBindingConfig>, +} + +impl KeymapConfig { + /// Returns true if no keybinding overrides are configured in any mode. + pub(crate) fn is_empty(&self) -> bool { + self.emacs.is_empty() + && self.vim_normal.is_empty() + && self.vim_insert.is_empty() + && self.inspector.is_empty() + && self.prefix.is_empty() + } +} + +#[derive(Clone, Debug, Deserialize, Serialize)] +pub(crate) struct Preview { + pub(crate) strategy: PreviewStrategy, +} + +#[derive(Clone, Debug, Deserialize, Serialize)] +pub struct Daemon { + /// The daemon will handle sync on an interval. How often to sync, in seconds. + pub sync_frequency: u64, + + /// The path to the unix socket used by the daemon + pub socket_path: String, + + /// Path to the daemon pidfile used for process coordination. + pub pidfile_path: String, + + /// Use a socket passed via systemd's socket activation protocol, instead of the path + pub systemd_socket: bool, + + /// The port that should be used for TCP on non unix systems + pub tcp_port: u64, +} + +#[derive(Clone, Debug, Deserialize, Serialize)] +pub(crate) struct Search { + /// The list of enabled filter modes, in order of priority. + pub(crate) filters: Vec<FilterMode>, + + /// The recency score multiplier for the search index (default: 1.0). + /// Values < 1.0 reduce weight, > 1.0 increase weight, 0.0 disables. + pub(crate) recency_score_multiplier: f64, + + /// The frequency score multiplier for the search index (default: 1.0). + /// Values < 1.0 reduce weight, > 1.0 increase weight, 0.0 disables. + pub(crate) frequency_score_multiplier: f64, + + /// The overall frecency score multiplier for the search index (default: 1.0). + /// Applied after combining recency and frequency scores. + pub(crate) frecency_score_multiplier: f64, +} + +/// Log level for file logging. Maps to tracing's [`LevelFilter`]. +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Deserialize, Serialize)] +#[serde(rename_all = "lowercase")] +pub(crate) enum LogLevel { + Trace, + Debug, + #[default] + Info, + Warn, + Error, +} + +impl LogLevel { + /// Convert to a tracing directive string for use with [`EnvFilter`]. + pub(crate) fn as_directive(self) -> &'static str { + match self { + Self::Trace => "trace", + Self::Debug => "debug", + Self::Info => "info", + Self::Warn => "warn", + Self::Error => "error", + } + } +} + +/// Configuration for a specific log type (search or daemon). +#[derive(Clone, Debug, Default, Deserialize, Serialize)] +pub(crate) struct LogConfig { + /// Log file name (relative to dir) or absolute path. + pub(crate) file: String, + + /// Override global enabled setting for this log type. + pub(crate) enabled: Option<bool>, + + /// Override global level setting for this log type. + pub(crate) level: Option<LogLevel>, + + /// Override global retention days setting for this log type. + pub(crate) retention: Option<u64>, +} + +#[derive(Clone, Debug, Deserialize, Serialize)] +pub(crate) struct Logs { + /// Enable file logging globally. Defaults to true. + #[serde(default = "Logs::default_enabled")] + pub(crate) enabled: bool, + + /// Directory for log files. Defaults to ~/.atuin/logs + pub(crate) dir: String, + + /// Default log level for file logging. Defaults to "info". + /// Note: [`ATUIN_LOG`] environment variable overrides this. + #[serde(default)] + pub(crate) level: LogLevel, + + /// Default retention days for log files. Defaults to 4. + #[serde(default = "Logs::default_retention")] + pub(crate) retention: u64, + + /// Search log settings + #[serde(default)] + pub(crate) search: LogConfig, + + /// Daemon log settings + #[serde(default)] + pub(crate) daemon: LogConfig, +} + +impl Default for Preview { + fn default() -> Self { + Self { + strategy: PreviewStrategy::Auto, + } + } +} + +impl Default for Daemon { + fn default() -> Self { + Self { + sync_frequency: 300, + socket_path: String::new(), + pidfile_path: String::new(), + systemd_socket: false, + tcp_port: 8889, + } + } +} + +impl Default for Logs { + fn default() -> Self { + Self { + enabled: true, + dir: String::new(), + level: LogLevel::default(), + retention: Self::default_retention(), + search: LogConfig { + file: "search.log".to_string(), + ..Default::default() + }, + daemon: LogConfig { + file: "daemon.log".to_string(), + ..Default::default() + }, + } + } +} + +impl Logs { + fn default_enabled() -> bool { + true + } + + fn default_retention() -> u64 { + 4 + } + + /// Returns whether search logging is enabled. + /// Uses search-specific setting if set, otherwise falls back to global. + pub(crate) fn search_enabled(&self) -> bool { + self.search.enabled.unwrap_or(self.enabled) + } + + /// Returns whether daemon logging is enabled. + /// Uses daemon-specific setting if set, otherwise falls back to global. + pub(crate) fn daemon_enabled(&self) -> bool { + self.daemon.enabled.unwrap_or(self.enabled) + } + + /// Returns the log level for search logging. + /// Uses search-specific setting if set, otherwise falls back to global. + pub(crate) fn search_level(&self) -> LogLevel { + self.search.level.unwrap_or(self.level) + } + + /// Returns the log level for daemon logging. + /// Uses daemon-specific setting if set, otherwise falls back to global. + pub(crate) fn daemon_level(&self) -> LogLevel { + self.daemon.level.unwrap_or(self.level) + } + + /// Returns the retention days for search logging. + /// Uses search-specific setting if set, otherwise falls back to global. + pub(crate) fn search_retention(&self) -> u64 { + self.search.retention.unwrap_or(self.retention) + } + + /// Returns the retention days for daemon logging. + /// Uses daemon-specific setting if set, otherwise falls back to global. + pub(crate) fn daemon_retention(&self) -> u64 { + self.daemon.retention.unwrap_or(self.retention) + } +} + +impl Default for Search { + fn default() -> Self { + Self { + filters: vec![ + FilterMode::Global, + FilterMode::Host, + FilterMode::Session, + FilterMode::SessionPreload, + FilterMode::Workspace, + FilterMode::Directory, + ], + + recency_score_multiplier: 1.0, + frequency_score_multiplier: 1.0, + frecency_score_multiplier: 1.0, + } + } +} + +// The preview height strategy also takes max_preview_height into account. +#[derive(Clone, Debug, Deserialize, Copy, PartialEq, Eq, ValueEnum, Serialize)] +pub(crate) enum PreviewStrategy { + // Preview height is calculated for the length of the selected command. + #[serde(rename = "auto")] + Auto, + + // Preview height is calculated for the length of the longest command stored in the history. + #[serde(rename = "static")] + Static, + + // max_preview_height is used as fixed height. + #[serde(rename = "fixed")] + Fixed, +} + +/// Column types available for the interactive search UI. +#[derive(Clone, Copy, Debug, Deserialize, PartialEq, Eq, Serialize)] +#[serde(rename_all = "lowercase")] +pub(crate) enum UiColumnType { + /// Command execution duration (e.g., "123ms") + Duration, + /// Relative time since execution (e.g., "59s ago") + Time, + /// Absolute timestamp (e.g., "2025-01-22 14:35") + Datetime, + /// Working directory + Directory, + /// Hostname + Host, + /// Username + User, + /// Exit code + Exit, + /// The command itself (should be last, expands to fill) + Command, +} + +impl UiColumnType { + /// Returns the default width for this column type (in characters). + /// The Command column returns 0 as it expands to fill remaining space. + pub(crate) fn default_width(self) -> u16 { + match self { + Self::Duration => 5, // "814ms" + Self::Time => 9, // "459ms ago" + Self::Datetime => 16, // "2025-01-22 14:35" + Self::Directory => 20, + Self::Host => 15, + Self::User => 10, + Self::Exit => { + if cfg!(windows) { + 11 // 32-bit integer on Windows: "-1978335212" + } else { + 3 // Usually a byte on Unix + } + } + Self::Command => 0, // Expands to fill + } + } +} + +/// A column configuration with type and optional custom width. +/// Can be specified as just a string (uses default width) or as an object with type and width. +#[derive(Clone, Debug, Serialize)] +pub(crate) struct UiColumn { + pub(crate) column_type: UiColumnType, + pub(crate) width: u16, + /// If true, this column expands to fill remaining space. Only one column should expand. + pub(crate) expand: bool, +} + +impl UiColumn { + pub(crate) fn new(column_type: UiColumnType) -> Self { + Self { + width: column_type.default_width(), + expand: column_type == UiColumnType::Command, + column_type, + } + } +} + +// Custom deserialize to handle both string and object formats: +// "duration" or { type = "duration", width = 8, expand = true } +impl<'de> Deserialize<'de> for UiColumn { + fn deserialize<D>(deserializer: D) -> Result<Self, D::Error> + where + D: serde::Deserializer<'de>, + { + use serde::de::{self, MapAccess, Visitor}; + + struct UiColumnVisitor; + + impl<'de> Visitor<'de> for UiColumnVisitor { + type Value = UiColumn; + + fn expecting(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.write_str( + "a column type string or an object with 'type' and optional 'width'/'expand'", + ) + } + + fn visit_str<E>(self, value: &str) -> Result<UiColumn, E> + where + E: de::Error, + { + let column_type: UiColumnType = + Deserialize::deserialize(de::value::StrDeserializer::new(value))?; + Ok(UiColumn::new(column_type)) + } + + fn visit_map<M>(self, mut map: M) -> Result<UiColumn, M::Error> + where + M: MapAccess<'de>, + { + let mut column_type: Option<UiColumnType> = None; + let mut width: Option<u16> = None; + let mut expand: Option<bool> = None; + + while let Some(key) = map.next_key::<String>()? { + match key.as_str() { + "type" => { + column_type = Some(map.next_value()?); + } + "width" => { + width = Some(map.next_value()?); + } + "expand" => { + expand = Some(map.next_value()?); + } + _ => { + let _: de::IgnoredAny = map.next_value()?; + } + } + } + + let column_type = column_type.ok_or_else(|| de::Error::missing_field("type"))?; + let width = width.unwrap_or_else(|| column_type.default_width()); + let expand = expand.unwrap_or(column_type == UiColumnType::Command); + Ok(UiColumn { + column_type, + width, + expand, + }) + } + } + + deserializer.deserialize_any(UiColumnVisitor) + } +} + +/// UI-specific settings for the interactive search. +#[derive(Clone, Debug, Deserialize, Serialize)] +pub(crate) struct Ui { + /// Columns to display in interactive search, from left to right. + /// The indicator column (" > ") is always shown first implicitly. + /// The "command" column should be last as it expands to fill remaining space. + /// Can be simple strings or objects with type and width. + #[serde(default = "Ui::default_columns")] + pub(crate) columns: Vec<UiColumn>, +} + +impl Ui { + fn default_columns() -> Vec<UiColumn> { + vec![ + UiColumn::new(UiColumnType::Duration), + UiColumn::new(UiColumnType::Time), + UiColumn::new(UiColumnType::Command), + ] + } + + /// Validate the UI configuration. + /// Returns an error if more than one column has expand = true. + pub(crate) fn validate(&self) -> Result<()> { + let expand_count = self.columns.iter().filter(|c| c.expand).count(); + if expand_count > 1 { + bail!( + "Only one column can have expand = true, but {} columns are set to expand", + expand_count + ); + } + Ok(()) + } +} + +impl Default for Ui { + fn default() -> Self { + Self { + columns: Self::default_columns(), + } + } +} + +/// Sync-specific settings. +#[derive(Clone, Debug, Default, Deserialize, Serialize)] +pub(crate) struct Sync { + /// The sync address for atuin. + pub(crate) address: String, + + #[serde(default)] + pub(crate) frequency: String, + + #[serde(default)] + pub(crate) auto: bool, + + #[serde(default)] + pub(crate) user_id_path: Option<PathBuf>, + + #[serde(default)] + pub(crate) encryption_key_path: Option<PathBuf>, +} + +impl Sync { + fn try_read_file(file: Option<&PathBuf>) -> Result<Option<String>> { + if let Some(path) = file { + if path.try_exists()? { + let user = read_to_string(path)?; + + if user.is_empty() { + Ok(None) + } else { + Ok(Some(user)) + } + } else { + // It's okay that the file doesn't exist. + // The important part is to error out if we can't access it (e.g. Because of missing + // permissions). + Ok(None) + } + } else { + Ok(None) + } + } + + pub(crate) fn have_sync_user(&self) -> Result<bool> { + let sa = self.user_id()?; + Ok(sa.is_some()) + } + + pub(crate) fn user_id(&self) -> Result<Option<Uuid>> { + Self::try_read_file(self.user_id_path.as_ref())? + .map(|file| { + Uuid::parse_str(file.trim()).context( + "Failed to decode user id as UUID, while trying to decode sync user_id", + ) + }) + .transpose() + } + pub(crate) fn encryption_key(&self) -> Result<Option<Key>> { + Self::try_read_file(self.encryption_key_path.as_ref())? + .as_deref() + .map(str::trim) + .map(decode_key) + .transpose() + } +} + +#[derive(Clone, Debug, Deserialize, Serialize)] +#[expect(clippy::struct_excessive_bools)] +pub struct Settings { + pub(crate) data_dir: Option<String>, + pub(crate) dialect: Dialect, + pub(crate) timezone: Timezone, + pub(crate) style: Style, + + pub db_path: String, + pub record_store_path: String, + pub(crate) search_mode: SearchMode, + pub(crate) filter_mode: Option<FilterMode>, + pub(crate) filter_mode_shell_up_key_binding: Option<FilterMode>, + pub(crate) search_mode_shell_up_key_binding: Option<SearchMode>, + pub(crate) shell_up_key_binding: bool, + pub(crate) inline_height: u16, + pub(crate) inline_height_shell_up_key_binding: Option<u16>, + pub(crate) invert: bool, + pub(crate) show_preview: bool, + pub(crate) max_preview_height: u16, + pub(crate) show_help: bool, + pub(crate) show_tabs: bool, + pub(crate) show_numeric_shortcuts: bool, + pub(crate) auto_hide_height: u16, + pub(crate) exit_mode: ExitMode, + pub(crate) keymap_mode: KeymapMode, + pub(crate) keymap_mode_shell: KeymapMode, + pub(crate) keymap_cursor: HashMap<String, CursorStyle>, + pub(crate) word_jump_mode: WordJumpMode, + pub(crate) word_chars: String, + pub(crate) scroll_context_lines: usize, + pub(crate) history_format: String, + pub(crate) strip_trailing_whitespace: bool, + pub(crate) prefers_reduced_motion: bool, + pub(crate) store_failed: bool, + pub(crate) no_mouse: bool, + + #[serde(with = "serde_regex", default = "RegexSet::empty", skip_serializing)] + pub(crate) history_filter: RegexSet, + + #[serde(with = "serde_regex", default = "RegexSet::empty", skip_serializing)] + pub(crate) cwd_filter: RegexSet, + + pub(crate) secrets_filter: bool, + pub(crate) workspaces: bool, + pub(crate) ctrl_n_shortcuts: bool, + + pub(crate) network_connect_timeout: u64, + pub(crate) network_timeout: u64, + pub local_timeout: f64, + pub(crate) enter_accept: bool, + pub(crate) smart_sort: bool, + pub(crate) command_chaining: bool, + + #[serde(default)] + pub(crate) sync: Sync, + + #[serde(default)] + pub(crate) stats: Stats, + + #[serde(default)] + pub(crate) keys: Keys, + + #[serde(default)] + pub(crate) keymap: KeymapConfig, + + #[serde(default)] + pub(crate) preview: Preview, + + #[serde(default)] + pub daemon: Daemon, + + #[serde(default)] + pub(crate) search: Search, + + #[serde(default)] + pub(crate) ui: Ui, + + #[serde(default)] + pub(crate) logs: Logs, + + #[serde(default)] + pub(crate) meta: meta::Settings, +} + +impl Settings { + // -- Meta store: lazily initialized on first access -- + + pub(crate) async fn meta_store() -> Result<&'static crate::aclient::meta::MetaStore> { + META_STORE + .get_or_try_init(|| async { + let (db_path, timeout) = META_CONFIG.get().ok_or_else(|| { + eyre!("meta store config not set — Settings::new() has not been called") + })?; + crate::aclient::meta::MetaStore::new(db_path, *timeout).await + }) + .await + } + + pub(crate) async fn host_id() -> Result<HostId> { + Self::meta_store().await?.host_id().await + } + + pub(crate) async fn last_sync() -> Result<OffsetDateTime> { + Self::meta_store().await?.last_sync().await + } + + pub(crate) async fn save_sync_time() -> Result<()> { + Self::meta_store().await?.save_sync_time().await + } + + pub(crate) fn default_filter_mode(&self, git_root: bool) -> FilterMode { + self.filter_mode + .filter(|x| self.search.filters.contains(x)) + .or_else(|| { + self.search + .filters + .iter() + .find(|x| match (x, git_root, self.workspaces) { + (FilterMode::Workspace, true, true) => true, + (FilterMode::Workspace, _, _) => false, + (_, _, _) => true, + }) + .copied() + }) + .unwrap_or(FilterMode::Global) + } + + pub(crate) fn builder() -> Result<ConfigBuilder<DefaultState>> { + Self::builder_with_data_dir(&utils::data_dir()) + } + + #[expect(clippy::too_many_lines)] + fn builder_with_data_dir(data_dir: &std::path::Path) -> Result<ConfigBuilder<DefaultState>> { + let db_path = data_dir.join("history.db"); + let record_store_path = data_dir.join("records.db"); + let kv_path = data_dir.join("kv.db"); + let scripts_path = data_dir.join("scripts.db"); + let ai_sessions_path = data_dir.join("ai_sessions.db"); + let socket_path = utils::runtime_dir().join("atuin.sock"); + let pidfile_path = data_dir.join("atuin-daemon.pid"); + let logs_dir = utils::logs_dir(); + + let key_path = data_dir.join("key"); + let meta_path = data_dir.join("meta.db"); + + Ok(Config::builder() + .set_default("history_format", "{time}\t{command}\t{duration}")? + .set_default("db_path", db_path.to_str())? + .set_default("record_store_path", record_store_path.to_str())? + .set_default("key_path", key_path.to_str())? + .set_default("dialect", "us")? + .set_default("timezone", "local")? + .set_default("auto_sync", true)? + .set_default("sync.address", "https://api.atuin.sh")? + .set_default("sync_frequency", "5m")? + .set_default("search_mode", "fuzzy")? + .set_default("filter_mode", None::<String>)? + .set_default("style", "compact")? + .set_default("inline_height", 40)? + .set_default("show_preview", true)? + .set_default("preview.strategy", "auto")? + .set_default("max_preview_height", 4)? + .set_default("show_help", true)? + .set_default("show_tabs", true)? + .set_default("show_numeric_shortcuts", true)? + .set_default("auto_hide_height", 8)? + .set_default("invert", false)? + .set_default("exit_mode", "return-original")? + .set_default("word_jump_mode", "emacs")? + .set_default( + "word_chars", + "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789", + )? + .set_default("scroll_context_lines", 1)? + .set_default("shell_up_key_binding", false)? + .set_default("workspaces", false)? + .set_default("ctrl_n_shortcuts", false)? + .set_default("secrets_filter", true)? + .set_default("strip_trailing_whitespace", true)? + .set_default("network_connect_timeout", 5)? + .set_default("network_timeout", 30)? + .set_default("local_timeout", 2.0)? + // enter_accept defaults to false here, but true in the default config file. The dissonance is + // intentional! + // Existing users will get the default "False", so we don't mess with any potential + // muscle memory. + // New users will get the new default, that is more similar to what they are used to. + .set_default("enter_accept", false)? + .set_default("keys.scroll_exits", true)? + .set_default("keys.accept_past_line_end", true)? + .set_default("keys.exit_past_line_start", true)? + .set_default("keys.accept_past_line_start", false)? + .set_default("keys.accept_with_backspace", false)? + .set_default("keys.prefix", "a")? + .set_default("keymap_mode", "emacs")? + .set_default("keymap_mode_shell", "auto")? + .set_default("keymap_cursor", HashMap::<String, String>::new())? + .set_default("smart_sort", false)? + .set_default("command_chaining", false)? + .set_default("store_failed", true)? + .set_default("daemon.sync_frequency", 300)? + .set_default("daemon.socket_path", socket_path.to_str())? + .set_default("daemon.pidfile_path", pidfile_path.to_str())? + .set_default("daemon.systemd_socket", false)? + .set_default("daemon.tcp_port", 8889)? + .set_default("logs.enabled", true)? + .set_default("logs.dir", logs_dir.to_str())? + .set_default("logs.level", "info")? + .set_default("logs.search.file", "search.log")? + .set_default("logs.daemon.file", "daemon.log")? + .set_default("logs.ai.file", "ai.log")? + .set_default("kv.db_path", kv_path.to_str())? + .set_default("scripts.db_path", scripts_path.to_str())? + .set_default("search.recency_score_multiplier", 1.0)? + .set_default("search.frequency_score_multiplier", 1.0)? + .set_default("search.frecency_score_multiplier", 1.0)? + .set_default("meta.db_path", meta_path.to_str())? + .set_default("ai.db_path", ai_sessions_path.to_str())? + .set_default("ai.session_continue_minutes", 60)? + .set_default("ai.send_cwd", false)? + .set_default("ai.opening.send_cwd", false)? + .set_default("ai.opening.send_last_command", false)? + .set_default( + "search.filters", + vec![ + "global", + "host", + "session", + "workspace", + "directory", + "session-preload", + ], + )? + .set_default("theme.name", "default")? + .set_default("theme.debug", None::<bool>)? + .set_default("tmux.enabled", false)? + .set_default("tmux.width", "80%")? + .set_default("tmux.height", "60%")? + .set_default( + "prefers_reduced_motion", + std::env::var("NO_MOTION").ok().map_or_else( + || config::Value::new(None, config::ValueKind::Boolean(false)), + |_| config::Value::new(None, config::ValueKind::Boolean(true)), + ), + )? + .set_default("no_mouse", false)? + .add_source( + Environment::with_prefix("atuin") + .prefix_separator("_") + .separator("__"), + )) + } + + pub(crate) fn get_config_path() -> Result<PathBuf> { + let config_dir = utils::config_dir(); + + create_dir_all(&config_dir) + .wrap_err_with(|| format!("could not create dir {}", config_dir.display()))?; + + let mut config_file = std::env::var("ATUIN_CONFIG_DIR").map_or_else( + |_| { + let mut config_file = PathBuf::new(); + config_file.push(config_dir); + config_file + }, + PathBuf::from, + ); + + config_file.push("config.toml"); + + Ok(config_file) + } + + /// Build a merged `Config` from defaults, config file, and environment. + /// + /// This resolves `data_dir`, initializes the data directory on disk, + /// and layers defaults → config file → env overrides. Both `new()` and + /// `get_config_value()` use this so the resolution logic lives in one place. + fn build_config() -> Result<Config> { + let config_file = Self::get_config_path()?; + + // extract data_dir first so we can use it as the base for other path defaults + let effective_data_dir = if config_file.exists() { + #[derive(Deserialize, Default)] + struct DataDirOnly { + data_dir: Option<String>, + } + + let config_file_str = config_file + .to_str() + .ok_or_else(|| eyre!("config file path is not valid UTF-8"))?; + + let partial_config = Config::builder() + .add_source(ConfigFile::new(config_file_str, FileFormat::Toml)) + .add_source( + Environment::with_prefix("atuin") + .prefix_separator("_") + .separator("__"), + ) + .build() + .ok(); + + let custom_data_dir = partial_config + .and_then(|c| c.try_deserialize::<DataDirOnly>().ok()) + .and_then(|d| d.data_dir); + + match custom_data_dir { + Some(dir) => { + let expanded = shellexpand::full(&dir) + .map_err(|e| eyre!("failed to expand data_dir path: {}", e))?; + PathBuf::from(expanded.as_ref()) + } + None => utils::data_dir(), + } + } else { + utils::data_dir() + }; + + DATA_DIR.set(effective_data_dir.clone()).ok(); + + create_dir_all(&effective_data_dir) + .wrap_err_with(|| format!("could not create dir {}", effective_data_dir.display()))?; + + let mut config_builder = Self::builder_with_data_dir(&effective_data_dir)?; + + config_builder = if config_file.exists() { + let config_file_str = config_file + .to_str() + .ok_or_else(|| eyre!("config file path is not valid UTF-8"))?; + config_builder.add_source(ConfigFile::new(config_file_str, FileFormat::Toml)) + } else { + // TODO(@bpeetz): Rework the config handling, so that we can actually auto-write a + // file with defaults. <2026-06-13> + create_dir_all(config_file.parent().unwrap())?; + + info!( + "No config file at: `{}`. Not adding one.", + config_file.display() + ); + + config_builder + }; + + // all paths should be expanded + let built = config_builder.build_cloned()?; + config_builder = [ + "db_path", + "record_store_path", + "key_path", + "daemon.socket_path", + "daemon.pidfile_path", + "logs.dir", + "logs.search.file", + "logs.daemon.file", + ] + .iter() + .map(|key| (key, built.get_string(key).unwrap_or_default())) + .filter_map(|(key, value)| match Self::expand_path(&value) { + Ok(expanded) => Some((key, expanded)), + Err(e) => { + log::warn!("failed to expand path for {key}: {e}"); + None + } + }) + .fold(config_builder, |builder, (key, value)| { + builder + .set_override(key, value) + .unwrap_or_else(|_| panic!("failed to set absolute path override for {key}")) + }); + + config_builder.build().map_err(Into::into) + } + + /// Look up a single config value by dotted key (e.g. `"daemon.sync_frequency"`). + /// + /// Returns the effective value after merging defaults, config file, and + /// environment — without the side-effects of full `Settings` construction + /// (meta store init, path expansion, etc.). + pub(crate) fn get_config_value(key: &str) -> Result<String> { + let config = Self::build_config()?; + let value: config::Value = config + .get(key) + .map_err(|e| eyre!("failed to get config value '{}': {}", key, e))?; + Ok(Self::format_resolved_value(&value, key)) + } + + fn format_resolved_value(value: &config::Value, prefix: &str) -> String { + use config::ValueKind; + + match &value.kind { + ValueKind::Nil => String::new(), + ValueKind::Boolean(b) => b.to_string(), + ValueKind::I64(i) => i.to_string(), + ValueKind::I128(i) => i.to_string(), + ValueKind::U64(u) => u.to_string(), + ValueKind::U128(u) => u.to_string(), + ValueKind::Float(f) => f.to_string(), + ValueKind::String(s) => s.clone(), + ValueKind::Array(arr) => { + let items: Vec<String> = arr + .iter() + .map(|v| Self::format_resolved_value(v, "")) + .collect(); + format!("[{}]", items.join(", ")) + } + ValueKind::Table(map) => { + let mut lines = Vec::new(); + let mut keys: Vec<_> = map.keys().collect(); + keys.sort(); + + for k in keys { + let v = &map[k]; + let full_key = if prefix.is_empty() { + k.clone() + } else { + format!("{prefix}.{k}") + }; + + match &v.kind { + ValueKind::Table(_) => { + lines.push(Self::format_resolved_value(v, &full_key)); + } + _ => { + lines.push(format!( + "{} = {}", + full_key, + Self::format_resolved_value(v, "") + )); + } + } + } + + lines.join("\n") + } + } + } + + pub fn new() -> Result<Self> { + let config = Self::build_config()?; + let settings: Self = config + .try_deserialize() + .map_err(|e| eyre!("failed to deserialize: {}", e))?; + + // Validate UI settings + settings.ui.validate()?; + + // Register meta store config for lazy initialization on first access + META_CONFIG + .set((settings.meta.db_path.clone(), settings.local_timeout)) + .ok(); + + Ok(settings) + } + + fn expand_path(path: &str) -> Result<String> { + shellexpand::full(&path) + .map(|p| p.to_string()) + .map_err(|e| eyre!("failed to expand path: {}", e)) + } + + pub(crate) fn paths_ok(&self) -> bool { + // TODO(@bpeetz): Add the `sync.*` paths <2026-06-11> + let paths = [&self.db_path, &self.record_store_path, &self.meta.db_path]; + paths.iter().all(|p| !utils::broken_symlink(p)) + } +} + +impl Default for Settings { + fn default() -> Self { + // if this panics something is very wrong, as the default config + // does not build or deserialize into the settings struct + Self::builder() + .expect("Could not build default") + .build() + .expect("Could not build config") + .try_deserialize() + .expect("Could not deserialize config") + } +} + +#[cfg(test)] +pub(crate) fn test_local_timeout() -> f64 { + std::env::var("ATUIN_TEST_LOCAL_TIMEOUT") + .ok() + .and_then(|x| x.parse().ok()) + // this hardcoded value should be replaced by a simple way to get the + // default local_timeout of Settings if possible + .unwrap_or(2.0) +} + +#[cfg(test)] +mod tests { + use std::str::FromStr; + + use eyre::Result; + + use super::Timezone; + + #[test] + fn can_parse_offset_timezone_spec() -> Result<()> { + assert_eq!(Timezone::from_str("+02")?.0.as_hms(), (2, 0, 0)); + assert_eq!(Timezone::from_str("-04")?.0.as_hms(), (-4, 0, 0)); + assert_eq!(Timezone::from_str("+05:30")?.0.as_hms(), (5, 30, 0)); + assert_eq!(Timezone::from_str("-09:30")?.0.as_hms(), (-9, -30, 0)); + + // single digit hours are allowed + assert_eq!(Timezone::from_str("+2")?.0.as_hms(), (2, 0, 0)); + assert_eq!(Timezone::from_str("-4")?.0.as_hms(), (-4, 0, 0)); + assert_eq!(Timezone::from_str("+5:30")?.0.as_hms(), (5, 30, 0)); + assert_eq!(Timezone::from_str("-9:30")?.0.as_hms(), (-9, -30, 0)); + + // fully qualified form + assert_eq!(Timezone::from_str("+09:30:00")?.0.as_hms(), (9, 30, 0)); + assert_eq!(Timezone::from_str("-09:30:00")?.0.as_hms(), (-9, -30, 0)); + + // these offsets don't really exist but are supported anyway + assert_eq!(Timezone::from_str("+0:5")?.0.as_hms(), (0, 5, 0)); + assert_eq!(Timezone::from_str("-0:5")?.0.as_hms(), (0, -5, 0)); + assert_eq!(Timezone::from_str("+01:23:45")?.0.as_hms(), (1, 23, 45)); + assert_eq!(Timezone::from_str("-01:23:45")?.0.as_hms(), (-1, -23, -45)); + + // require a leading sign for clarity + assert!(Timezone::from_str("5").is_err()); + assert!(Timezone::from_str("10:30").is_err()); + + Ok(()) + } + + #[test] + fn can_choose_workspace_filters_when_in_git_context() -> Result<()> { + let mut settings = super::Settings::default(); + settings.search.filters = vec![ + super::FilterMode::Workspace, + super::FilterMode::Host, + super::FilterMode::Directory, + super::FilterMode::Session, + super::FilterMode::Global, + ]; + settings.workspaces = true; + + assert_eq!( + settings.default_filter_mode(true), + super::FilterMode::Workspace, + ); + + Ok(()) + } + + #[test] + fn wont_choose_workspace_filters_when_not_in_git_context() -> Result<()> { + let mut settings = super::Settings::default(); + settings.search.filters = vec![ + super::FilterMode::Workspace, + super::FilterMode::Host, + super::FilterMode::Directory, + super::FilterMode::Session, + super::FilterMode::Global, + ]; + settings.workspaces = true; + + assert_eq!(settings.default_filter_mode(false), super::FilterMode::Host,); + + Ok(()) + } + + #[test] + fn wont_choose_workspace_filters_when_workspaces_disabled() -> Result<()> { + let mut settings = super::Settings::default(); + settings.search.filters = vec![ + super::FilterMode::Workspace, + super::FilterMode::Host, + super::FilterMode::Directory, + super::FilterMode::Session, + super::FilterMode::Global, + ]; + settings.workspaces = false; + + assert_eq!(settings.default_filter_mode(true), super::FilterMode::Host,); + + Ok(()) + } + + #[test] + fn builder_with_data_dir_uses_custom_paths() -> Result<()> { + use std::path::PathBuf; + + let custom_dir = PathBuf::from("/custom/data/dir"); + let builder = super::Settings::builder_with_data_dir(&custom_dir)?; + let config = builder.build()?; + + let db_path: String = config.get("db_path")?; + let key_path: String = config.get("key_path")?; + let record_store_path: String = config.get("record_store_path")?; + let kv_db_path: String = config.get("kv.db_path")?; + let scripts_db_path: String = config.get("scripts.db_path")?; + let meta_db_path: String = config.get("meta.db_path")?; + let daemon_socket_path: String = config.get("daemon.socket_path")?; + let daemon_pidfile_path: String = config.get("daemon.pidfile_path")?; + + assert_eq!(db_path, custom_dir.join("history.db").to_str().unwrap()); + assert_eq!(key_path, custom_dir.join("key").to_str().unwrap()); + assert_eq!( + record_store_path, + custom_dir.join("records.db").to_str().unwrap() + ); + assert_eq!(kv_db_path, custom_dir.join("kv.db").to_str().unwrap()); + assert_eq!( + scripts_db_path, + custom_dir.join("scripts.db").to_str().unwrap() + ); + assert_eq!(meta_db_path, custom_dir.join("meta.db").to_str().unwrap()); + assert_eq!( + daemon_socket_path, + turtle_common::utils::runtime_dir() + .join("atuin.sock") + .to_str() + .unwrap() + ); + assert_eq!( + daemon_pidfile_path, + custom_dir.join("atuin-daemon.pid").to_str().unwrap() + ); + + Ok(()) + } + + #[test] + fn keymap_config_deserializes_simple_binding() { + let json = r#"{"emacs": {"ctrl-c": "exit"}}"#; + let config: super::KeymapConfig = serde_json::from_str(json).unwrap(); + assert_eq!(config.emacs.len(), 1); + match &config.emacs["ctrl-c"] { + super::KeyBindingConfig::Simple(s) => assert_eq!(s, "exit"), + _ => panic!("expected Simple variant"), + } + } + + #[test] + fn keymap_config_deserializes_conditional_binding() { + let json = r#"{ + "emacs": { + "left": [ + {"when": "cursor-at-start", "action": "exit"}, + {"action": "cursor-left"} + ] + } + }"#; + let config: super::KeymapConfig = serde_json::from_str(json).unwrap(); + match &config.emacs["left"] { + super::KeyBindingConfig::Rules(rules) => { + assert_eq!(rules.len(), 2); + assert_eq!(rules[0].when.as_deref(), Some("cursor-at-start")); + assert_eq!(rules[0].action, "exit"); + assert!(rules[1].when.is_none()); + assert_eq!(rules[1].action, "cursor-left"); + } + _ => panic!("expected Rules variant"), + } + } + + #[test] + fn keymap_config_deserializes_vim_normal() { + let json = r#"{"vim-normal": {"j": "select-next", "k": "select-previous"}}"#; + let config: super::KeymapConfig = serde_json::from_str(json).unwrap(); + assert_eq!(config.vim_normal.len(), 2); + assert!(config.emacs.is_empty()); + } + + #[test] + fn keymap_config_is_empty_when_default() { + let config = super::KeymapConfig::default(); + assert!(config.is_empty()); + } + + #[test] + fn keymap_config_mixed_modes() { + let json = r#"{ + "emacs": {"ctrl-c": "exit"}, + "vim-normal": {"q": "exit"}, + "inspector": {"d": "delete"} + }"#; + let config: super::KeymapConfig = serde_json::from_str(json).unwrap(); + assert!(!config.is_empty()); + assert_eq!(config.emacs.len(), 1); + assert_eq!(config.vim_normal.len(), 1); + assert_eq!(config.inspector.len(), 1); + assert!(config.vim_insert.is_empty()); + assert!(config.prefix.is_empty()); + } +} diff --git a/crates/daemon/src/aclient/settings/watcher.rs b/crates/daemon/src/aclient/settings/watcher.rs new file mode 100644 index 00000000..01d20855 --- /dev/null +++ b/crates/daemon/src/aclient/settings/watcher.rs @@ -0,0 +1,260 @@ +//! Config file watching for automatic settings reload. +//! +//! This module provides a `SettingsWatcher` that monitors the config file +//! for changes and broadcasts updated settings via a `tokio::sync::watch` channel. +//! +//! # Example +//! +//! ```no_run +//! use crate::aclient::settings::watcher::global_settings_watcher; +//! +//! async fn example() -> eyre::Result<()> { +//! let watcher = global_settings_watcher()?; +//! let mut rx = watcher.subscribe(); +//! +//! // React to settings changes +//! while rx.changed().await.is_ok() { +//! let settings = rx.borrow(); +//! println!("Settings updated!"); +//! } +//! Ok(()) +//! } +//! ``` + +use std::{ + path::{Path, PathBuf}, + sync::{Arc, OnceLock}, + time::Duration, +}; + +use eyre::{Result, WrapErr}; +use log::{debug, error, info, warn}; +use notify::{ + Config as NotifyConfig, RecommendedWatcher, RecursiveMode, Watcher, + event::{EventKind, ModifyKind}, +}; +use tokio::sync::watch; + +use super::Settings; + +/// Global singleton for the settings watcher. +static SETTINGS_WATCHER: OnceLock<Result<SettingsWatcher, String>> = OnceLock::new(); + +/// Get the global settings watcher singleton. +/// +/// Initializes the watcher on first call. Subsequent calls return the same instance. +/// The watcher monitors the config file for changes and broadcasts updates. +pub(crate) fn global_settings_watcher() -> Result<&'static SettingsWatcher> { + let result = SETTINGS_WATCHER.get_or_init(|| SettingsWatcher::new().map_err(|e| e.to_string())); + + match result { + Ok(watcher) => Ok(watcher), + Err(e) => Err(eyre::eyre!("{}", e)), + } +} + +/// Watches the config file for changes and broadcasts updated settings. +/// +/// Uses `notify` for cross-platform file watching and `tokio::sync::watch` +/// for efficient broadcast to multiple subscribers. +pub(crate) struct SettingsWatcher { + /// Receiver for settings updates. Clone this to subscribe. + rx: watch::Receiver<Arc<Settings>>, + /// Keeps the file watcher alive for the lifetime of this struct. + _watcher: RecommendedWatcher, +} + +impl SettingsWatcher { + /// Create a new settings watcher. + /// + /// Loads initial settings and starts watching the config file for changes. + /// Changes are debounced (500ms) to avoid multiple reloads during saves. + pub(crate) fn new() -> Result<Self> { + let initial_settings = Arc::new(Settings::new()?); + let (tx, rx) = watch::channel(initial_settings); + + let config_path = Self::config_path(); + info!("starting config file watcher: {}", config_path.display()); + + let watcher = Self::create_watcher(tx, &config_path)?; + + Ok(Self { + rx, + _watcher: watcher, + }) + } + + /// Subscribe to settings updates. + /// + /// Returns a receiver that will be notified when settings change. + /// Use `changed().await` to wait for the next update, then `borrow()` + /// to access the current settings. + pub(crate) fn subscribe(&self) -> watch::Receiver<Arc<Settings>> { + self.rx.clone() + } + + /// Get the config file path. + fn config_path() -> PathBuf { + let config_dir = std::env::var("ATUIN_CONFIG_DIR") + .map_or_else(|_| turtle_common::utils::config_dir(), PathBuf::from); + config_dir.join("config.toml") + } + + /// Create the file watcher with debouncing. + fn create_watcher( + tx: watch::Sender<Arc<Settings>>, + config_path: &Path, + ) -> Result<RecommendedWatcher> { + // Channel for debouncing file events + let (debounce_tx, debounce_rx) = std::sync::mpsc::channel::<()>(); + + // Spawn debounce thread + let config_path_clone = config_path.to_owned(); + std::thread::spawn(move || { + Self::debounce_loop(&debounce_rx, &tx, &config_path_clone); + }); + + // Clone config_path for use in the watcher callback + let config_path_for_watcher = config_path.to_owned(); + + // Canonicalize config path for reliable comparison on macOS + // (handles symlinks like /var -> /private/var) + let canonical_config_path = config_path_for_watcher + .canonicalize() + .unwrap_or_else(|_| config_path_for_watcher.clone()); + + // Create file watcher + let mut watcher = RecommendedWatcher::new( + move |res: Result<notify::Event, notify::Error>| { + match res { + Ok(event) => { + // Defensive: if paths is empty, we can't filter, so assume + // it might be our config file and trigger a reload to be safe + if event.paths.is_empty() { + warn!( + "config watcher: event has no paths, triggering reload to be safe" + ); + debounce_tx.send(()).expect("should still be active"); + return; + } + + // Only react to events for our specific config file + // (filter out editor temp files, backups, etc.) + let is_config_file = event.paths.iter().any(|path| { + // Canonicalize for reliable comparison (handles macOS symlinks) + let canonical_event_path = + path.canonicalize().unwrap_or_else(|_| path.clone()); + + // Check if this event is for our config file + // (either exact match or the file was renamed to our config) + canonical_event_path == canonical_config_path + || path.file_name() == config_path_for_watcher.file_name() + }); + + if !is_config_file { + return; + } + + // Only react to modify events (content changes) or creates + if matches!( + event.kind, + EventKind::Modify(ModifyKind::Data(_) | ModifyKind::Any) + | EventKind::Create(_) + ) { + debug!("config file event detected: {event:?}"); + // Send to debounce channel (ignore send errors - receiver might be gone) + debounce_tx.send(()).ok(); + } + } + Err(e) => { + error!("file watcher error: {e}"); + } + } + }, + NotifyConfig::default(), + ) + .wrap_err("failed to create file watcher")?; + + // Watch the config file's parent directory (some editors create new files) + let watch_path = config_path.parent().unwrap_or(config_path); + + // Defensive: ensure watch path exists before trying to watch + if !watch_path.exists() { + warn!( + "config directory does not exist, creating it: {}", + watch_path.display() + ); + std::fs::create_dir_all(watch_path).wrap_err_with(|| { + format!( + "failed to create config directory: {}", + watch_path.display() + ) + })?; + } + + watcher + .watch(watch_path, RecursiveMode::NonRecursive) + .wrap_err_with(|| { + format!("failed to watch config directory: {}", watch_path.display()) + })?; + + info!( + "config file watcher initialized for: {}", + watch_path.display() + ); + Ok(watcher) + } + + /// Debounce loop that batches file events and reloads settings. + fn debounce_loop( + rx: &std::sync::mpsc::Receiver<()>, + tx: &watch::Sender<Arc<Settings>>, + config_path: &Path, + ) { + const DEBOUNCE_DURATION: Duration = Duration::from_millis(500); + + loop { + // Wait for first event + if rx.recv().is_err() { + // Channel closed, watcher was dropped + debug!("config watcher debounce loop exiting"); + return; + } + + // Drain any additional events within debounce window + while rx.recv_timeout(DEBOUNCE_DURATION).is_ok() { + // Keep draining + } + + // Defensive: check if config file exists before reloading + // (handles case where file was deleted - we'll get notified when it's recreated) + if !config_path.exists() { + debug!( + "config file does not exist, skipping reload: {}", + config_path.display() + ); + continue; + } + + // Now reload settings + info!( + "config file changed, reloading settings: {}", + config_path.display() + ); + match Settings::new() { + Ok(settings) => { + if tx.send(Arc::new(settings)).is_err() { + // All receivers dropped + debug!("all settings subscribers dropped, exiting"); + return; + } + info!("settings reloaded successfully"); + } + Err(e) => { + warn!("failed to reload settings: {e}"); + // Keep the old settings, don't broadcast the error + } + } + } + } +} diff --git a/crates/daemon/src/aclient/utils.rs b/crates/daemon/src/aclient/utils.rs new file mode 100644 index 00000000..cf515183 --- /dev/null +++ b/crates/daemon/src/aclient/utils.rs @@ -0,0 +1,92 @@ +pub(crate) fn get_hostname() -> String { + std::env::var("ATUIN_HOST_NAME") + .unwrap_or_else(|_| whoami::hostname().unwrap_or_else(|_| "unknown-host".to_string())) +} + +pub(crate) fn get_username() -> String { + std::env::var("ATUIN_HOST_USER") + .unwrap_or_else(|_| whoami::username().unwrap_or_else(|_| "unknown-user".to_string())) +} + +/// Returns a pair of the hostname and username, separated by a colon. +pub(crate) fn get_host_user() -> String { + format!("{}:{}", get_hostname(), get_username()) +} + +/// Setup a [`SQLite`] database. +/// +/// This takes care of correct locking, so that we avoid a race when setting up the database. +macro_rules! setup_db { + ( + $db_path:expr, + $a_timeout:expr, + $opts:expr, + $m_name:literal $(,)? + ) => {{ + async fn migrate(pool: &SqlitePool) -> sqlx::Result<()> { + { sqlx::sqlx_macros::migrate!($m_name) }.run(pool).await?; + Ok(()) + } + + crate::aclient::utils::setup_db_inner($db_path, $a_timeout, $opts, migrate) + }}; +} +pub(crate) use setup_db; + +use std::{os::fd::AsRawFd, path::Path, time::Duration}; + +use fs_err::OpenOptions; +use sqlx::{ + SqlitePool, + sqlite::{SqliteConnectOptions, SqlitePoolOptions}, +}; +use tracing::debug; + +/// Helper for `setup_db!` +pub(crate) async fn setup_db_inner( + db_path: &Path, + acquire_timeout: f64, + mk_opts: fn(&str) -> sqlx::Result<SqliteConnectOptions>, + migrate: impl AsyncFn(&SqlitePool) -> sqlx::Result<()>, +) -> sqlx::Result<SqlitePool> { + async fn open_db(timeout: f64, opts: SqliteConnectOptions) -> sqlx::Result<SqlitePool> { + let pool = SqlitePoolOptions::new() + .acquire_timeout(Duration::from_secs_f64(timeout)) + .connect_with(opts) + .await?; + + Ok(pool) + } + + { + let file = OpenOptions::new() + .read(true) + .write(true) + .create(true) + .open(db_path)?; + + // Lock the db file while we are running the migrations. + // Why? Because there is a small chance that we start running migrations (e.g. as the daemon) + // and then another process is started, which will also try to run migrations. + // Essentially, one of the processes will receive with a SQLite UNIQUE constraint failure. + // So let's avoid that possibility from the start. + file.lock()?; + + let pool = open_db( + acquire_timeout, + mk_opts(format!("/proc/self/fd/{}", file.as_raw_fd()).as_str())?, + ) + .await?; + + debug!("running sqlite database setup"); + + migrate(&pool).await?; + + file.unlock()?; + } + + let real_opts = mk_opts(db_path.to_str().expect("Should be utf-8"))?; + let pool = open_db(acquire_timeout, real_opts).await?; + + Ok(pool) +} diff --git a/crates/daemon/src/client.rs b/crates/daemon/src/client.rs index 5cccb5ff..5f8ea0f2 100644 --- a/crates/daemon/src/client.rs +++ b/crates/daemon/src/client.rs @@ -7,36 +7,25 @@ use hyper_util::rt::TokioIo; #[cfg(unix)] use tokio::net::UnixStream; -use tracing::{Level, instrument, span}; +use crate::aclient::{history::History, settings::Settings}; use crate::generated; use crate::{ - atuin_client::{ - database::Context, - history::History, - settings::{FilterMode, Settings}, - }, - atuin_daemon::{ - events::DaemonEvent, - generated::{ - control::{ - ForceSyncEvent, HistoryDeletedEvent, HistoryPrunedEvent, HistoryRebuiltEvent, - SendEventRequest, SettingsReloadedEvent, ShutdownEvent, - control_client::ControlClient as ControlServiceClient, - }, - history::{ - EndHistoryReply, EndHistoryRequest, ShutdownRequest, StartHistoryReply, - StartHistoryRequest, StatusReply, StatusRequest, TailHistoryReply, - TailHistoryRequest, history_client::HistoryClient as HistoryServiceClient, - }, - search::{ - FilterMode as RpcFilterMode, SearchContext as RpcSearchContext, SearchRequest, - SearchResponse, search_client::SearchClient as SearchServiceClient, - }, - semantic::{ - CommandCapture, RecordCommandsReply, - semantic_client::SemanticClient as SemanticServiceClient, - }, + events::DaemonEvent, + generated::{ + control::{ + ForceSyncEvent, HistoryDeletedEvent, HistoryPrunedEvent, HistoryRebuiltEvent, + SendEventRequest, SettingsReloadedEvent, ShutdownEvent, + control_client::ControlClient as ControlServiceClient, + }, + history::{ + EndHistoryReply, EndHistoryRequest, ShutdownRequest, StartHistoryReply, + StartHistoryRequest, StatusReply, StatusRequest, TailHistoryReply, TailHistoryRequest, + history_client::HistoryClient as HistoryServiceClient, + }, + semantic::{ + CommandCapture, RecordCommandsReply, + semantic_client::SemanticClient as SemanticServiceClient, }, }, }; @@ -129,7 +118,7 @@ impl HistoryClient { Ok(self.client.status(StatusRequest {}).await?.into_inner()) } - pub async fn tail_history(&mut self) -> Result<tonic::Streaming<TailHistoryReply>> { + pub(crate) async fn tail_history(&mut self) -> Result<tonic::Streaming<TailHistoryReply>> { Ok(self .client .tail_history(TailHistoryRequest {}) @@ -143,92 +132,13 @@ impl HistoryClient { } } -pub struct SearchClient { - client: SearchServiceClient<Channel>, -} - -impl SearchClient { - #[cfg(unix)] - pub async fn new(path: String) -> Result<Self> { - let log_path = path.clone(); - let channel = Endpoint::try_from("http://atuin_local_daemon:0")? - .connect_with_connector(service_fn(move |_: Uri| { - let path = path.clone(); - - async move { - Ok::<_, std::io::Error>(TokioIo::new(UnixStream::connect(path.clone()).await?)) - } - })) - .await - .wrap_err_with(|| { - format!( - "failed to connect to local atuin daemon at {}. Is it running?", - &log_path - ) - })?; - - let client = SearchServiceClient::new(channel); - - Ok(Self { client }) - } - - #[instrument(skip_all, level = Level::TRACE, name = "daemon_client_search", fields(query = %query, query_id = query_id))] - pub async fn search( - &mut self, - query: String, - query_id: u64, - filter_mode: FilterMode, - context: Option<Context>, - ) -> Result<tonic::Streaming<SearchResponse>> { - let request = SearchRequest { - query, - query_id, - filter_mode: RpcFilterMode::from(filter_mode).into(), - context: context.map(RpcSearchContext::from), - }; - let request_stream = tokio_stream::once(request); - let response = span!(Level::TRACE, "daemon_client_search.request") - .in_scope(async || self.client.search(request_stream).await) - .await?; - - Ok(response.into_inner()) - } -} - -impl From<FilterMode> for RpcFilterMode { - fn from(filter_mode: FilterMode) -> Self { - match filter_mode { - FilterMode::Global => Self::Global, - FilterMode::Host => Self::Host, - FilterMode::Session => Self::Session, - FilterMode::Directory => Self::Directory, - FilterMode::Workspace => Self::Workspace, - FilterMode::SessionPreload => Self::SessionPreload, - } - } -} - -impl From<Context> for RpcSearchContext { - fn from(context: Context) -> Self { - Self { - session_id: context.session, - cwd: context.cwd, - hostname: context.hostname, - host_id: context.host_id, - git_root: context - .git_root - .map(|path| path.to_string_lossy().to_string()), - } - } -} - -pub struct SemanticClient { +pub(crate) struct SemanticClient { client: SemanticServiceClient<Channel>, } impl SemanticClient { #[cfg(unix)] - pub async fn new(path: String) -> Result<Self> { + pub(crate) async fn new(path: String) -> Result<Self> { let log_path = path.clone(); let channel = Endpoint::try_from("http://atuin_local_daemon:0")? .connect_with_connector(service_fn(move |_: Uri| { @@ -252,11 +162,11 @@ impl SemanticClient { } #[cfg(unix)] - pub async fn from_settings(settings: &Settings) -> Result<Self> { + pub(crate) async fn from_settings(settings: &Settings) -> Result<Self> { Self::new(settings.daemon.socket_path.clone()).await } - pub async fn record_commands( + pub(crate) async fn record_commands( &mut self, captures: Vec<CommandCapture>, ) -> Result<RecordCommandsReply> { @@ -279,7 +189,7 @@ pub struct ControlClient { impl ControlClient { /// Connect to the daemon's control service. #[cfg(unix)] - pub async fn new(path: String) -> Result<Self> { + pub(crate) async fn new(path: String) -> Result<Self> { let log_path = path.clone(); let channel = Endpoint::try_from("http://atuin_local_daemon:0")? .connect_with_connector(service_fn(move |_: Uri| { diff --git a/crates/daemon/src/components/history.rs b/crates/daemon/src/components/history.rs index a75ff774..7a0882a3 100644 --- a/crates/daemon/src/components/history.rs +++ b/crates/daemon/src/components/history.rs @@ -4,7 +4,7 @@ use std::{pin::Pin, sync::Arc}; -use crate::atuin_client::{ +use crate::aclient::{ history::{History, HistoryId, store::HistoryStore}, settings::Settings, }; @@ -35,7 +35,7 @@ const DAEMON_PROTOCOL_VERSION: u32 = 1; /// - Saves completed commands to the database and record store /// - Emits history events for other components (e.g., search indexing) /// - Provides the History gRPC service -pub struct HistoryComponent { +pub(crate) struct HistoryComponent { inner: Arc<HistoryComponentInner>, } @@ -52,7 +52,7 @@ struct HistoryComponentInner { impl HistoryComponent { /// Create a new history component. - pub fn new() -> Self { + pub(crate) fn new() -> Self { Self { inner: Arc::new(HistoryComponentInner { running: DashMap::new(), @@ -65,7 +65,7 @@ impl HistoryComponent { /// Get the gRPC service for this component. /// /// This returns a tonic service that can be added to a gRPC server. - pub fn grpc_service(&self) -> HistoryServer<HistoryGrpcService> { + pub(crate) fn grpc_service(&self) -> HistoryServer<HistoryGrpcService> { HistoryServer::new(HistoryGrpcService { inner: self.inner.clone(), }) @@ -111,7 +111,7 @@ impl Component for HistoryComponent { /// The gRPC service implementation. /// /// This is a thin wrapper that delegates to the component's shared state. -pub struct HistoryGrpcService { +pub(crate) struct HistoryGrpcService { inner: Arc<HistoryComponentInner>, } diff --git a/crates/daemon/src/components/mod.rs b/crates/daemon/src/components/mod.rs index 447e31df..7c478efb 100644 --- a/crates/daemon/src/components/mod.rs +++ b/crates/daemon/src/components/mod.rs @@ -10,16 +10,13 @@ //! Available components: //! //! - [`history::HistoryComponent`]: Command history lifecycle management -//! - [`search::SearchComponent`]: Fuzzy search over history //! - [`semantic::SemanticComponent`]: In-memory semantic command captures //! - [`sync::SyncComponent`]: Cloud sync -pub mod history; -pub mod search; -pub mod semantic; -pub mod sync; +pub(crate) mod history; +pub(crate) mod semantic; +pub(crate) mod sync; -pub use history::HistoryComponent; -pub use search::SearchComponent; -pub use semantic::SemanticComponent; -pub use sync::SyncComponent; +pub(crate) use history::HistoryComponent; +pub(crate) use semantic::SemanticComponent; +pub(crate) use sync::SyncComponent; diff --git a/crates/daemon/src/components/search.rs b/crates/daemon/src/components/search.rs deleted file mode 100644 index 91f2db17..00000000 --- a/crates/daemon/src/components/search.rs +++ /dev/null @@ -1,407 +0,0 @@ -//! Search component. -//! -//! Provides fuzzy search over command history using the Nucleo search library -//! with frecency-based ranking and dynamic filtering. - -use std::{pin::Pin, sync::Arc}; - -use eyre::Result; -use tokio::sync::RwLock; -use tokio_stream::Stream; -use tonic::{Request, Response, Status, Streaming}; -use tracing::{Level, debug, info, instrument, span, trace}; -use uuid::Uuid; - -use crate::{ - daemon::{Component, DaemonHandle}, - events::DaemonEvent, - generated::search::{ - self, FilterMode, SearchRequest, SearchResponse, - search_server::{Search as SearchSvc, SearchServer}, - }, - search::{IndexFilterMode, QueryContext, SearchIndex}, -}; - -const PAGE_SIZE: usize = 5000; -const RESULTS_LIMIT: u32 = 200; -/// How often to rebuild the frecency map (in seconds). -const FRECENCY_REFRESH_INTERVAL_SECS: u64 = 60; - -/// Search component - provides fuzzy search over command history. -/// -/// This component: -/// - Maintains a deduplicated search index with frecency ranking -/// - Loads history from the database on startup -/// - Updates the index when history events occur -/// - Provides the Search gRPC service -pub struct SearchComponent { - index: Arc<RwLock<SearchIndex>>, - handle: RwLock<Option<DaemonHandle>>, - loader_handle: Option<tokio::task::JoinHandle<()>>, - frecency_handle: Option<tokio::task::JoinHandle<()>>, -} - -impl SearchComponent { - /// Create a new search component. - pub fn new() -> Self { - Self { - index: Arc::new(RwLock::new(SearchIndex::new())), - handle: RwLock::new(None), - loader_handle: None, - frecency_handle: None, - } - } - - /// Get the gRPC service for this component. - pub fn grpc_service(&self) -> SearchServer<SearchGrpcService> { - SearchServer::new(SearchGrpcService { - index: self.index.clone(), - }) - } - - /// Rebuild the entire search index from the database. - #[expect(clippy::significant_drop_tightening, reason = "false positive")] - async fn rebuild_index(&self) -> Result<()> { - let handle_guard = self.handle.read().await; - let handle = handle_guard - .as_ref() - .ok_or_else(|| eyre::eyre!("component not initialized"))?; - - info!("Rebuilding search index from database"); - - // Create a new index - let new_index = SearchIndex::new(); - - // Load all history into the new index - let db = handle.history_db().clone(); - let mut pager = db.all_paged(PAGE_SIZE, false, true); - loop { - match pager.next().await { - Ok(Some(histories)) => { - info!( - "Loading {} history entries into search index", - histories.len() - ); - new_index.add_histories(&histories); - } - Ok(None) => break, - Err(e) => { - tracing::error!("Failed to load history during rebuild: {}", e); - break; - } - } - } - - info!( - "Search index rebuild complete; {} unique commands", - new_index.command_count() - ); - - // Replace the old index with the new one - *self.index.write().await = new_index; - Ok(()) - } -} - -impl Default for SearchComponent { - fn default() -> Self { - Self::new() - } -} - -#[tonic::async_trait] -impl Component for SearchComponent { - fn name(&self) -> &'static str { - "search" - } - - #[expect(clippy::significant_drop_tightening, reason = "false positive")] - async fn start(&mut self, handle: DaemonHandle) -> Result<()> { - *self.handle.write().await = Some(handle.clone()); - - // Spawn background task to load history into index - let index = self.index.clone(); - let db = handle.history_db().clone(); - let handle_for_loader = handle.clone(); - - self.loader_handle = Some(tokio::spawn(async move { - info!( - "Loading history into search index; page size = {}", - PAGE_SIZE - ); - let mut pager = db.all_paged(PAGE_SIZE, false, true); - loop { - match pager.next().await { - Ok(Some(histories)) => { - info!( - "Loading {} history entries into search index", - histories.len() - ); - index.read().await.add_histories(&histories); - } - Ok(None) => { - info!( - "Initial history load complete; {} unique commands indexed", - index.read().await.command_count() - ); - // Build initial frecency map with current settings - let settings = handle_for_loader.settings().await; - index.read().await.rebuild_frecency(&settings.search).await; - info!("Initial frecency map built"); - break; - } - Err(e) => { - tracing::error!("Failed to load history: {}", e); - break; - } - } - } - })); - - // Spawn background task to periodically refresh frecency - let index_for_frecency = self.index.clone(); - let handle_for_frecency = handle.clone(); - self.frecency_handle = Some(tokio::spawn(async move { - let mut interval = tokio::time::interval(std::time::Duration::from_secs( - FRECENCY_REFRESH_INTERVAL_SECS, - )); - loop { - interval.tick().await; - trace!("Refreshing frecency map"); - let settings = handle_for_frecency.settings().await; - index_for_frecency - .read() - .await - .rebuild_frecency(&settings.search) - .await; - } - })); - - tracing::info!("search component started"); - Ok(()) - } - - #[expect(clippy::significant_drop_tightening, reason = "false positive")] - async fn handle_event(&mut self, event: &DaemonEvent) -> Result<()> { - match event { - DaemonEvent::RecordsAdded(records) => { - debug!( - count = records.len(), - "Processing added records for search index" - ); - - let handle_guard = self.handle.read().await; - if let Some(handle) = handle_guard.as_ref() { - let histories: Vec<_> = handle - .history_db() - .query_history( - format!( - "select * from history where id in ({})", - records - .iter() - .map(|record| record.0.to_string()) - .collect::<Vec<_>>() - .join(",") - ) - .as_str(), - ) - .await - .unwrap_or_default(); - - span!(Level::TRACE, "inject_records", count = histories.len()) - .in_scope(async || { - self.index.read().await.add_histories(&histories); - }) - .await; - } - } - DaemonEvent::HistoryStarted(history) => { - debug!(id = %history.id, command = %history.command, "History started (no index action)"); - } - DaemonEvent::HistoryEnded(history) => { - span!(Level::TRACE, "inject_history_ended") - .in_scope(async || { - self.index.read().await.add_history(history); - }) - .await; - } - DaemonEvent::HistoryPruned | DaemonEvent::HistoryRebuilt => { - info!("History store pruned or rebuilt, rebuilding search index"); - if let Err(e) = self.rebuild_index().await { - tracing::error!("Failed to rebuild search index: {}", e); - } - } - DaemonEvent::HistoryDeleted { ids } => { - info!( - count = ids.len(), - "History deleted, rebuilding search index" - ); - // For now, just rebuild the entire index. A more efficient implementation - // would remove specific items from the index. - if let Err(e) = self.rebuild_index().await { - tracing::error!("Failed to rebuild search index: {}", e); - } - } - DaemonEvent::SettingsReloaded => { - info!("Settings reloaded, rebuilding frecency map with new multipliers"); - let handle_guard = self.handle.read().await; - if let Some(handle) = handle_guard.as_ref() { - let settings = handle.settings().await; - self.index - .read() - .await - .rebuild_frecency(&settings.search) - .await; - } - } - // Events we don't care about - DaemonEvent::SyncCompleted { .. } - | DaemonEvent::SyncFailed { .. } - | DaemonEvent::ForceSync - | DaemonEvent::ShutdownRequested => {} - } - Ok(()) - } - - async fn stop(&mut self) -> Result<()> { - if let Some(handle) = self.loader_handle.take() { - handle.abort(); - } - if let Some(handle) = self.frecency_handle.take() { - handle.abort(); - } - tracing::info!("search component stopped"); - Ok(()) - } -} - -/// The gRPC service implementation. -pub struct SearchGrpcService { - index: Arc<RwLock<SearchIndex>>, -} - -#[tonic::async_trait] -impl SearchSvc for SearchGrpcService { - type SearchStream = Pin<Box<dyn Stream<Item = Result<SearchResponse, Status>> + Send>>; - - #[instrument(skip_all, level = Level::TRACE, name = "search_rpc")] - async fn search( - &self, - request: Request<Streaming<SearchRequest>>, - ) -> Result<Response<Self::SearchStream>, Status> { - let mut in_stream = request.into_inner(); - let index = self.index.clone(); - - // Create output channel - let (tx, rx) = tokio::sync::mpsc::channel::<Result<SearchResponse, Status>>(128); - - // Spawn task to handle incoming requests and send responses - tokio::spawn(async move { - while let Some(req) = in_stream.message().await.transpose() { - match req { - Ok(search_req) => { - let query = search_req.query; - let query_id = search_req.query_id; - let filter_mode: FilterMode = search_req - .filter_mode - .try_into() - .unwrap_or(FilterMode::Global); - let proto_context = search_req.context; - - debug!( - "search request: query = {}, query_id = {}, filter_mode = {}, context = {:?}", - query, - query_id, - filter_mode.as_str_name(), - proto_context - ); - - // Convert proto FilterMode + context to IndexFilterMode - let index_filter = convert_filter_mode(filter_mode, proto_context.as_ref()); - - // Build QueryContext from proto context - let query_context = proto_context - .map(|ctx| QueryContext { - cwd: Some(with_trailing_slash(&ctx.cwd)), - git_root: ctx.git_root.map(|s| with_trailing_slash(&s)), - hostname: Some(ctx.hostname), - session_id: Some(ctx.session_id), - }) - .unwrap_or_default(); - - // Perform the search - let history_ids = - span!(Level::TRACE, "daemon_search_query", %query, query_id) - .in_scope(|| async { - let index = index.read().await; - index - .search(&query, index_filter, &query_context, RESULTS_LIMIT) - .await - }) - .await; - - // Convert history IDs to bytes - let ids: Vec<Vec<u8>> = history_ids - .iter() - .filter_map(|id| { - Uuid::parse_str(id) - .ok() - .map(|uuid| uuid.as_bytes().to_vec()) - }) - .collect(); - - if tx.send(Ok(SearchResponse { query_id, ids })).await.is_err() { - break; // Client disconnected - } - } - Err(e) => { - drop(tx.send(Err(e)).await); - break; - } - } - } - }); - - // Convert receiver to stream - let out_stream = tokio_stream::wrappers::ReceiverStream::new(rx); - Ok(Response::new(Box::pin(out_stream))) - } -} - -/// Convert proto `FilterMode` and context to `IndexFilterMode`. -fn convert_filter_mode( - mode: FilterMode, - context: Option<&search::SearchContext>, -) -> IndexFilterMode { - #[expect( - clippy::match_same_arms, - reason = "wildcard pattern used in second one" - )] - match (mode, context) { - (FilterMode::Global, _) => IndexFilterMode::Global, - (FilterMode::Directory, Some(ctx)) => { - IndexFilterMode::Directory(with_trailing_slash(&ctx.cwd)) - } - (FilterMode::Workspace, Some(ctx)) => ctx.git_root.as_ref().map_or_else( - || IndexFilterMode::Directory(with_trailing_slash(&ctx.cwd)), - |git_root| IndexFilterMode::Workspace(with_trailing_slash(git_root)), - ), - (FilterMode::Host, Some(ctx)) => IndexFilterMode::Host(ctx.hostname.clone()), - (FilterMode::Session, Some(ctx)) => IndexFilterMode::Session(ctx.session_id.clone()), - (FilterMode::SessionPreload, Some(ctx)) => { - // SessionPreload is similar to Session - filter by session - IndexFilterMode::Session(ctx.session_id.clone()) - } - // If no context provided, fall back to global - _ => IndexFilterMode::Global, - } -} - -#[cfg(not(windows))] -pub fn with_trailing_slash(s: &str) -> String { - if s.ends_with('/') { - s.to_string() - } else { - format!("{s}/") - } -} diff --git a/crates/daemon/src/components/semantic.rs b/crates/daemon/src/components/semantic.rs index 02f5c3d1..aec26887 100644 --- a/crates/daemon/src/components/semantic.rs +++ b/crates/daemon/src/components/semantic.rs @@ -8,7 +8,7 @@ use std::collections::{HashMap, VecDeque}; use std::fmt::{Display, Formatter}; use std::sync::Arc; -use crate::atuin_client::history::{History, HistoryId}; +use crate::aclient::history::{History, HistoryId}; use crate::generated::semantic; use eyre::Result; use tokio::sync::Mutex; @@ -30,7 +30,7 @@ const MAX_BYTES_PER_SESSION: usize = 32 * 1024 * 1024; const MAX_PENDING_HISTORIES: usize = 128; /// Stores completed command captures and associates them with history events. -pub struct SemanticComponent { +pub(crate) struct SemanticComponent { inner: Arc<SemanticComponentInner>, } @@ -84,7 +84,7 @@ struct SemanticCommandRecord { } impl SemanticComponent { - pub fn new() -> Self { + pub(crate) fn new() -> Self { Self { inner: Arc::new(SemanticComponentInner { state: Mutex::new(SemanticState::default()), @@ -92,7 +92,7 @@ impl SemanticComponent { } } - pub fn grpc_service(&self) -> SemanticServer<SemanticGrpcService> { + pub(crate) fn grpc_service(&self) -> SemanticServer<SemanticGrpcService> { SemanticServer::new(SemanticGrpcService { inner: self.inner.clone(), }) @@ -453,7 +453,7 @@ impl Display for SessionId { } } -pub struct SemanticGrpcService { +pub(crate) struct SemanticGrpcService { inner: Arc<SemanticComponentInner>, } diff --git a/crates/daemon/src/components/sync.rs b/crates/daemon/src/components/sync.rs index e898e8bd..e200ad73 100644 --- a/crates/daemon/src/components/sync.rs +++ b/crates/daemon/src/components/sync.rs @@ -9,7 +9,7 @@ use rand::Rng; use tokio::sync::mpsc; use tokio::time::{self, MissedTickBehavior}; -use crate::atuin_client::{history::store::HistoryStore, record::sync, settings::Settings}; +use crate::aclient::{history::store::HistoryStore, record::sync, settings::Settings}; use crate::{ daemon::{Component, DaemonHandle}, @@ -41,14 +41,14 @@ enum SyncState { /// - Implements exponential backoff on sync failures /// - Responds to [`ForceSync`] events for immediate sync /// - Emits SyncCompleted/SyncFailed events -pub struct SyncComponent { +pub(crate) struct SyncComponent { task_handle: Option<tokio::task::JoinHandle<()>>, command_tx: Option<mpsc::Sender<SyncCommand>>, } impl SyncComponent { /// Create a new sync component. - pub fn new() -> Self { + pub(crate) fn new() -> Self { Self { task_handle: None, command_tx: None, diff --git a/crates/daemon/src/control/mod.rs b/crates/daemon/src/control/mod.rs index fcc2a0b8..f727a3a2 100644 --- a/crates/daemon/src/control/mod.rs +++ b/crates/daemon/src/control/mod.rs @@ -7,15 +7,13 @@ use tonic::{Request, Response, Status}; use tracing::{Level, info, instrument}; use crate::{ - atuin_client::history::HistoryId, - atuin_daemon::{ - daemon::DaemonHandle, - events::DaemonEvent, - generated::control::{ - SendEventRequest, SendEventResponse, - control_server::{Control, ControlServer}, - send_event_request::Event, - }, + aclient::history::HistoryId, + daemon::DaemonHandle, + events::DaemonEvent, + generated::control::{ + SendEventRequest, SendEventResponse, + control_server::{Control, ControlServer}, + send_event_request::Event, }, }; @@ -23,18 +21,18 @@ use crate::{ /// /// This service is used by external processes to inject events into the daemon. /// It's not a component - it's part of the daemon's core infrastructure. -pub struct ControlService { +pub(crate) struct ControlService { handle: DaemonHandle, } impl ControlService { /// Create a new control service with the given daemon handle. - pub fn new(handle: DaemonHandle) -> Self { + pub(crate) fn new(handle: DaemonHandle) -> Self { Self { handle } } /// Get a tonic server for this service. - pub fn into_server(self) -> ControlServer<Self> { + pub(crate) fn into_server(self) -> ControlServer<Self> { ControlServer::new(self) } } diff --git a/crates/daemon/src/daemon.rs b/crates/daemon/src/daemon.rs index 8f0a5957..4e691be2 100644 --- a/crates/daemon/src/daemon.rs +++ b/crates/daemon/src/daemon.rs @@ -10,7 +10,7 @@ use std::sync::Arc; -use crate::atuin_client::{ +use crate::aclient::{ database::ClientSqlite as HistoryDatabase, encryption, record::sqlite_store::SqliteStore, settings::Settings, }; @@ -27,7 +27,7 @@ use crate::events::DaemonEvent; /// /// This contains all the resources that components and services need access to. /// The state is wrapped in an `Arc` and accessed via [`DaemonHandle`]. -pub struct DaemonState { +pub(crate) struct DaemonState { // Event bus event_tx: broadcast::Sender<DaemonEvent>, @@ -72,7 +72,7 @@ pub struct DaemonState { /// let history = handle.history_db().load(id).await?; /// ``` #[derive(Clone)] -pub struct DaemonHandle { +pub(crate) struct DaemonHandle { state: Arc<DaemonState>, } @@ -83,7 +83,7 @@ impl DaemonHandle { /// /// This is fire-and-forget - if no receivers are listening (which shouldn't /// happen in normal operation), the event is dropped silently. - pub fn emit(&self, event: DaemonEvent) { + pub(crate) fn emit(&self, event: DaemonEvent) { if let Err(e) = self.state.event_tx.send(event) { tracing::warn!("failed to emit event (no receivers?): {e}"); } @@ -94,12 +94,12 @@ impl DaemonHandle { /// Returns a receiver that will receive all events emitted after this call. /// Useful for components that need to listen for events outside of the /// normal `handle_event` callback flow. - pub fn subscribe(&self) -> broadcast::Receiver<DaemonEvent> { + pub(crate) fn subscribe(&self) -> broadcast::Receiver<DaemonEvent> { self.state.event_tx.subscribe() } /// Request graceful shutdown of the daemon. - pub fn shutdown(&self) { + pub(crate) fn shutdown(&self) { self.emit(DaemonEvent::ShutdownRequested); } @@ -109,7 +109,7 @@ impl DaemonHandle { /// /// This acquires a read lock on the settings. For most use cases, clone /// the settings if you need to hold onto them. - pub async fn settings(&self) -> tokio::sync::RwLockReadGuard<'_, Settings> { + pub(crate) async fn settings(&self) -> tokio::sync::RwLockReadGuard<'_, Settings> { self.state.settings.read().await } @@ -117,26 +117,26 @@ impl DaemonHandle { /// /// Use this when settings have already been loaded (e.g., from a file watcher) /// to avoid parsing the config file twice. - pub async fn apply_settings(&self, settings: Settings) { + pub(crate) async fn apply_settings(&self, settings: Settings) { *self.state.settings.write().await = settings; self.emit(DaemonEvent::SettingsReloaded); tracing::info!("settings applied"); } /// Get the encryption key. - pub fn encryption_key(&self) -> &[u8; 32] { + pub(crate) fn encryption_key(&self) -> &[u8; 32] { &self.state.encryption_key } // ---- Database ---- /// Get a reference to the history database. - pub fn history_db(&self) -> &HistoryDatabase { + pub(crate) fn history_db(&self) -> &HistoryDatabase { &self.state.history_db } /// Get a reference to the record store. - pub fn store(&self) -> &SqliteStore { + pub(crate) fn store(&self) -> &SqliteStore { &self.state.store } } @@ -171,7 +171,7 @@ impl std::fmt::Debug for DaemonHandle { /// # Example /// /// ```ignore -/// pub struct MyComponent { +/// pub(crate) struct MyComponent { /// handle: Option<DaemonHandle>, /// } /// @@ -203,7 +203,7 @@ impl std::fmt::Debug for DaemonHandle { /// } /// ``` #[tonic::async_trait] -pub trait Component: Send + Sync { +pub(crate) trait Component: Send + Sync { /// Human-readable name for logging and debugging. fn name(&self) -> &'static str; @@ -247,21 +247,21 @@ pub trait Component: Send + Sync { /// /// Events emitted during handling are queued and processed in subsequent /// iterations, ensuring the loop eventually drains. -pub struct Daemon { +pub(crate) struct Daemon { components: Vec<Box<dyn Component>>, handle: DaemonHandle, } impl Daemon { /// Create a new daemon builder. - pub fn builder(settings: Settings) -> DaemonBuilder { + pub(crate) fn builder(settings: Settings) -> DaemonBuilder { DaemonBuilder::new(settings) } /// Get a clone of the daemon handle. /// /// The handle can be used to emit events, access settings, etc. - pub fn handle(&self) -> DaemonHandle { + pub(crate) fn handle(&self) -> DaemonHandle { self.handle.clone() } @@ -269,7 +269,7 @@ impl Daemon { /// /// This must be called before `run_event_loop()`. It initializes all /// registered components with the daemon handle. - pub async fn start_components(&mut self) -> Result<()> { + pub(crate) async fn start_components(&mut self) -> Result<()> { for component in &mut self.components { tracing::info!(component = component.name(), "starting component"); component @@ -284,7 +284,7 @@ impl Daemon { /// /// This processes events until a [`ShutdownRequested`] event is received. /// Components must be started first via `start_components()`. - pub async fn run_event_loop(&mut self) -> Result<()> { + pub(crate) async fn run_event_loop(&mut self) -> Result<()> { let mut event_rx = self.handle.subscribe(); loop { match event_rx.recv().await { @@ -314,7 +314,7 @@ impl Daemon { /// Stop all components. /// /// This performs graceful shutdown of all components. - pub async fn stop_components(&mut self) { + pub(crate) async fn stop_components(&mut self) { for component in &mut self.components { tracing::info!(component = component.name(), "stopping component"); if let Err(e) = component.stop().await { @@ -361,7 +361,7 @@ impl Daemon { /// /// daemon.run().await?; /// ``` -pub struct DaemonBuilder { +pub(crate) struct DaemonBuilder { settings: Settings, store: Option<SqliteStore>, history_db: Option<HistoryDatabase>, @@ -370,7 +370,7 @@ pub struct DaemonBuilder { impl DaemonBuilder { /// Create a new daemon builder with the given settings. - pub fn new(settings: Settings) -> Self { + pub(crate) fn new(settings: Settings) -> Self { Self { settings, store: None, @@ -380,13 +380,13 @@ impl DaemonBuilder { } /// Set the record store. - pub fn store(mut self, store: SqliteStore) -> Self { + pub(crate) fn store(mut self, store: SqliteStore) -> Self { self.store = Some(store); self } /// Set the history database. - pub fn history_db(mut self, db: HistoryDatabase) -> Self { + pub(crate) fn history_db(mut self, db: HistoryDatabase) -> Self { self.history_db = Some(db); self } @@ -394,7 +394,7 @@ impl DaemonBuilder { /// Register a component. /// /// Components are started in registration order and stopped in reverse order. - pub fn component(mut self, component: impl Component + 'static) -> Self { + pub(crate) fn component(mut self, component: impl Component + 'static) -> Self { self.components.push(Box::new(component)); self } @@ -402,7 +402,7 @@ impl DaemonBuilder { /// Build the daemon. /// /// This loads the encryption key and creates the daemon state. - pub fn build(self) -> Result<Daemon> { + pub(crate) fn build(self) -> Result<Daemon> { let store = self.store.ok_or_else(|| eyre::eyre!("store is required"))?; let history_db = self .history_db diff --git a/crates/daemon/src/events.rs b/crates/daemon/src/events.rs index 32ed1ff1..092e5f32 100644 --- a/crates/daemon/src/events.rs +++ b/crates/daemon/src/events.rs @@ -7,7 +7,7 @@ //! External processes (like CLI commands) can also inject events via the //! Control gRPC service. -use crate::atuin_client::history::{History, HistoryId}; +use crate::aclient::history::{History, HistoryId}; use turtle_common::record::RecordId; /// Events that flow through the daemon's event bus. diff --git a/crates/daemon/src/generated.rs b/crates/daemon/src/generated.rs index 6620e94c..9deb4e0c 100644 --- a/crates/daemon/src/generated.rs +++ b/crates/daemon/src/generated.rs @@ -11,22 +11,14 @@ )] /// Semantic command capture gRPC service types. -pub mod semantic { +pub(crate) mod semantic { tonic::include_proto!("semantic"); } -/// Search module for the daemon gRPC search service. -/// -/// This module provides fuzzy search over command history using Nucleo. -pub mod search { - // Include the generated proto code - tonic::include_proto!("search"); -} - /// History module for the daemon gRPC history service. /// /// This module contains the proto-generated types for the history gRPC service. -pub mod history { +pub(crate) mod history { // Include the generated proto code tonic::include_proto!("history"); } @@ -35,10 +27,10 @@ pub mod history { /// /// This module provides the gRPC service that allows external processes /// (like CLI commands) to inject events into the daemon's event bus. -pub mod control { +pub(crate) mod control { // Include the generated proto code tonic::include_proto!("control"); // Re-export the service - pub use crate::control::ControlService; + pub(crate) use crate::control::ControlService; } diff --git a/crates/daemon/src/lib.rs b/crates/daemon/src/lib.rs index 1abf0314..a5d233c6 100644 --- a/crates/daemon/src/lib.rs +++ b/crates/daemon/src/lib.rs @@ -1,27 +1,25 @@ -use crate::atuin_client::database::ClientSqlite as HistoryDatabase; -use crate::atuin_client::record::sqlite_store::SqliteStore; -use crate::atuin_client::settings::{Settings, watcher::global_settings_watcher}; +use crate::aclient::database::ClientSqlite as HistoryDatabase; +use crate::aclient::record::sqlite_store::SqliteStore; +use crate::aclient::settings::{Settings, watcher::global_settings_watcher}; use eyre::Result; +pub mod aclient; pub mod client; -pub mod components; -pub mod control; -pub mod daemon; -pub mod events; -pub mod search; -pub mod server; -pub mod generated; +pub(crate) mod components; +pub(crate) mod control; +pub(crate) mod daemon; +pub(crate) mod events; +pub(crate) mod server; + +pub(crate) mod generated; // Re-export core daemon types for convenience -pub use daemon::Daemon; +pub(crate) use daemon::Daemon; pub use events::DaemonEvent; // Re-export components -pub use components::{HistoryComponent, SearchComponent, SemanticComponent, SyncComponent}; - -// Re-export client helpers -pub use client::SemanticClient; +pub(crate) use components::{HistoryComponent, SemanticComponent, SyncComponent}; /// Boot the daemon using the new component-based architecture. /// @@ -34,14 +32,12 @@ pub async fn boot( ) -> Result<()> { // Create the components let history_component = HistoryComponent::new(); - let search_component = SearchComponent::new(); let semantic_component = SemanticComponent::new(); let sync_component = SyncComponent::new(); // Get the gRPC services before moving components into the daemon // (The services share state with the components via Arc) let history_service = history_component.grpc_service(); - let search_service = search_component.grpc_service(); let semantic_service = semantic_component.grpc_service(); // Build the daemon @@ -49,7 +45,6 @@ pub async fn boot( .store(store) .history_db(history_db) .component(history_component) - .component(search_component) .component(semantic_component) .component(sync_component) .build()?; @@ -94,7 +89,6 @@ pub async fn boot( server::run_grpc_server( &settings, history_service, - search_service, semantic_service, control_service.into_server(), handle, diff --git a/crates/daemon/src/main.rs b/crates/daemon/src/main.rs new file mode 100644 index 00000000..26a5cafd --- /dev/null +++ b/crates/daemon/src/main.rs @@ -0,0 +1,377 @@ +#[allow(unused_imports)] +use clap::Parser; +use eyre::{Result, WrapErr, bail, eyre}; +use fs4::fs_std::FileExt; +use std::fs::{self, File, OpenOptions}; +use std::io::{ErrorKind, Write}; +#[cfg(unix)] +use std::os::unix::net::UnixStream as StdUnixStream; +use std::path::{Path, PathBuf}; +use std::process::{Command, Stdio}; +use std::time::{Duration, Instant}; +use tokio::time::sleep; +use turtle_daemon::{ + DaemonEvent, + aclient::{ + database::ClientSqlite, history::History, record::sqlite_store::SqliteStore, + settings::Settings, + }, + client::{ControlClient, DaemonClientErrorKind, HistoryClient, classify_error}, +}; + +#[derive(Parser, Debug)] +#[command(infer_subcommands = true)] +pub(crate) enum Cmd { + /// Start the daemon server + Start { + #[arg(long, hide = true)] + daemonize: bool, + + /// Also write daemon logs to the console (useful for debugging) + #[arg(long)] + show_logs: bool, + }, + + /// Show the daemon's current status + Status, + + /// Stop the daemon gracefully + Stop, +} + +impl Cmd { + pub(crate) async fn run( + self, + settings: Settings, + store: SqliteStore, + history_db: ClientSqlite, + ) -> Result<()> { + match self { + Cmd::Start { .. } => run(settings, store, history_db).await, + Cmd::Status => status_cmd(&settings).await, + Cmd::Stop => stop_cmd(&settings).await, + } + } +} + +#[tokio::main] +async fn main() -> Result<()> { + let settings = Settings::new().wrap_err("could not load client settings")?; + let db_path = PathBuf::from(settings.db_path.as_str()); + let record_store_path = PathBuf::from(settings.record_store_path.as_str()); + + let db = ClientSqlite::new(db_path, settings.local_timeout).await?; + let sqlite_store = SqliteStore::new(record_store_path, settings.local_timeout).await?; + + Cmd::parse().run(settings, sqlite_store, db).await +} + +const DAEMON_VERSION: &str = env!("CARGO_PKG_VERSION"); +const DAEMON_PROTOCOL_VERSION: u32 = 1; +const STARTUP_POLL: Duration = Duration::from_millis(40); +const LOCK_POLL: Duration = Duration::from_millis(20); +const LEGACY_DAEMON_RESTART_MESSAGE: &str = "legacy daemon detected; restart daemon manually"; + +struct PidfileGuard { + file: File, +} + +impl PidfileGuard { + fn acquire(path: &Path) -> Result<Self> { + let mut file = open_lock_file(path)?; + + if !file.try_lock_exclusive()? { + bail!( + "daemon already running (pidfile lock busy at {})", + path.display() + ); + } + + file.set_len(0) + .wrap_err_with(|| format!("could not truncate daemon pidfile {}", path.display()))?; + writeln!(file, "{}", std::process::id()) + .and_then(|()| writeln!(file, "{DAEMON_VERSION}")) + .wrap_err_with(|| format!("could not write daemon pidfile {}", path.display()))?; + + Ok(Self { file }) + } +} + +impl Drop for PidfileGuard { + fn drop(&mut self) { + drop(self.file.unlock()); + } +} + +enum Probe { + Ready(HistoryClient), + NeedsRestart(String), + Unreachable(eyre::Report), +} + +fn daemon_matches_expected(version: &str, protocol: u32) -> bool { + version == DAEMON_VERSION && protocol == DAEMON_PROTOCOL_VERSION +} + +fn daemon_mismatch_message(version: &str, protocol: u32) -> String { + if protocol == DAEMON_PROTOCOL_VERSION { + format!("daemon is out of date: expected {DAEMON_VERSION}, got {version}") + } else { + format!("daemon protocol mismatch: expected {DAEMON_PROTOCOL_VERSION}, got {protocol}") + } +} + +fn is_legacy_daemon_error(err: &eyre::Report) -> bool { + matches!(classify_error(err), DaemonClientErrorKind::Unimplemented) +} + +fn open_lock_file(path: &Path) -> Result<File> { + if let Some(parent) = path.parent() { + fs::create_dir_all(parent) + .wrap_err_with(|| format!("could not create lock directory {}", parent.display()))?; + } + + OpenOptions::new() + .read(true) + .write(true) + .create(true) + .truncate(false) + .open(path) + .wrap_err_with(|| format!("could not open lock file {}", path.display())) +} + +async fn wait_for_lock(path: &Path, timeout: Duration) -> Result<File> { + let file = open_lock_file(path)?; + let start = Instant::now(); + + loop { + match file.try_lock_exclusive() { + Ok(true) => return Ok(file), + Ok(false) => { + if start.elapsed() >= timeout { + bail!("timed out waiting for lock at {}", path.display()); + } + + sleep(LOCK_POLL).await; + } + Err(err) => { + return Err(eyre!("could not lock {}: {err}", path.display())); + } + } + } +} + +async fn wait_for_pidfile_available(path: &Path, timeout: Duration) -> Result<()> { + let file = wait_for_lock(path, timeout).await?; + file.unlock() + .wrap_err_with(|| format!("failed to unlock {}", path.display()))?; + Ok(()) +} + +async fn connect_client(settings: &Settings) -> Result<HistoryClient> { + HistoryClient::new( + #[cfg(unix)] + settings.daemon.socket_path.clone(), + ) + .await +} + +async fn probe(settings: &Settings) -> Probe { + let mut client = match connect_client(settings).await { + Ok(client) => client, + Err(err) => return Probe::Unreachable(err), + }; + + match client.status().await { + Ok(status) => { + if daemon_matches_expected(&status.version, status.protocol) { + Probe::Ready(client) + } else { + Probe::NeedsRestart(daemon_mismatch_message(&status.version, status.protocol)) + } + } + Err(err) => Probe::Unreachable(err), + } +} + +async fn request_shutdown(settings: &Settings) { + if let Ok(mut client) = connect_client(settings).await { + drop(client.shutdown().await); + } +} + +fn startup_timeout(settings: &Settings) -> Duration { + Duration::from_secs_f64(settings.local_timeout.max(0.5) + 2.0) +} + +pub(crate) async fn start_history(settings: &Settings, history: History) -> Result<String> { + match async { + connect_client(settings) + .await? + .start_history(history.clone()) + .await + } + .await + { + Ok(resp) => { + if daemon_matches_expected(&resp.version, resp.protocol) { + return Ok(resp.id); + } + + Err(eyre!( + "{}. Restart the daemon manually", + daemon_mismatch_message(&resp.version, resp.protocol) + )) + } + Err(err) => Err(err), + } +} + +pub(crate) async fn end_history( + settings: &Settings, + id: String, + duration: u64, + exit: i64, +) -> Result<()> { + match async { + connect_client(settings) + .await? + .end_history(id.clone(), duration, exit) + .await + } + .await + { + Ok(resp) => { + if daemon_matches_expected(&resp.version, resp.protocol) { + return Ok(()); + } + + Err(eyre!( + "{}. Restart the daemon manually", + daemon_mismatch_message(&resp.version, resp.protocol) + )) + } + Err(err) => Err(err), + } +} + +/// Emit a daemon event. +pub(crate) async fn emit_event(settings: &Settings, event: DaemonEvent) { + // Try to connect and send + match ControlClient::from_settings(settings).await { + Ok(mut client) => { + if let Err(e) = client.send_event(event).await { + tracing::debug!(?e, "failed to send event to daemon"); + } + } + Err(e) => { + tracing::debug!(?e, "daemon not available, skipping event emission"); + } + } +} + +pub(crate) async fn tail_client(settings: &Settings) -> Result<HistoryClient> { + match probe(settings).await { + Probe::Ready(client) => Ok(client), + Probe::NeedsRestart(reason) => { + bail!("{reason}. Restart the daemon manually"); + } + Probe::Unreachable(err) if is_legacy_daemon_error(&err) => { + Err(err.wrap_err(LEGACY_DAEMON_RESTART_MESSAGE)) + } + Probe::Unreachable(err) => Err(err), + } +} + +async fn status_cmd(settings: &Settings) -> Result<()> { + match probe(settings).await { + Probe::Ready(mut client) => { + let status = client.status().await?; + println!("Daemon running"); + println!(" PID: {}", status.pid); + println!(" Version: {}", status.version); + println!(" Protocol: {}", status.protocol); + println!(" Healthy: {}", status.healthy); + #[cfg(unix)] + println!(" Socket: {}", settings.daemon.socket_path); + } + Probe::NeedsRestart(reason) => { + println!("Daemon running (needs restart)"); + println!(" Reason: {reason}"); + } + Probe::Unreachable(_) => { + println!("Daemon is not running"); + } + } + + Ok(()) +} + +async fn stop_cmd(settings: &Settings) -> Result<()> { + let Ok(mut client) = connect_client(settings).await else { + println!("Daemon is not running"); + return Ok(()); + }; + + match client.shutdown().await { + Ok(true) => { + println!("Shutdown requested"); + + let pidfile_path = PathBuf::from(&settings.daemon.pidfile_path); + let timeout = Duration::from_secs(5); + match wait_for_pidfile_available(&pidfile_path, timeout).await { + Ok(()) => println!("Daemon stopped"), + Err(_) => println!("Daemon may still be shutting down"), + } + + Ok(()) + } + Ok(false) => bail!("Daemon rejected shutdown request"), + Err(err) => Err(err.wrap_err("Failed to send shutdown request")), + } +} + +async fn run(settings: Settings, store: SqliteStore, history_db: ClientSqlite) -> Result<()> { + let pidfile_path = PathBuf::from(&settings.daemon.pidfile_path); + let _pidfile_guard = PidfileGuard::acquire(&pidfile_path)?; + + turtle_daemon::boot(settings, store, history_db).await?; + + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::{ + DAEMON_PROTOCOL_VERSION, DAEMON_VERSION, daemon_matches_expected, daemon_mismatch_message, + }; + + #[test] + fn test_version_matches() { + assert!(daemon_matches_expected( + DAEMON_VERSION, + DAEMON_PROTOCOL_VERSION + )); + } + + #[test] + fn test_version_mismatch() { + assert!(!daemon_matches_expected("0.0.0", DAEMON_PROTOCOL_VERSION)); + assert!(!daemon_matches_expected(DAEMON_VERSION, 999)); + assert!(!daemon_matches_expected("0.0.0", 999)); + } + + #[test] + fn test_mismatch_message_version() { + let msg = daemon_mismatch_message("0.0.0", DAEMON_PROTOCOL_VERSION); + assert!(msg.contains("out of date"), "got: {msg}"); + assert!(msg.contains("0.0.0")); + assert!(msg.contains(DAEMON_VERSION)); + } + + #[test] + fn test_mismatch_message_protocol() { + let msg = daemon_mismatch_message(DAEMON_VERSION, 999); + assert!(msg.contains("protocol mismatch"), "got: {msg}"); + } +} diff --git a/crates/daemon/src/search/mod.rs b/crates/daemon/src/search/mod.rs deleted file mode 100644 index b4d03bcd..00000000 --- a/crates/daemon/src/search/mod.rs +++ /dev/null @@ -1,557 +0,0 @@ -//! Search index with frecency-based ranking. -//! -//! This module provides a deduplicated search index where each unique command -//! is stored once, with metadata about all its invocations. This enables: -//! -//! - Efficient fuzzy matching (fewer items to match) -//! - Frecency-based ranking (frequency + recency) -//! - Dynamic filtering by directory, host, session, etc. - -use std::{ - collections::{HashMap, HashSet}, - sync::Arc, -}; - -use crate::atuin_client::settings::Search; -use crate::{ - atuin_client::history::History, atuin_daemon::components::search::with_trailing_slash, -}; -use atuin_nucleo::{Injector, Nucleo, pattern}; -use dashmap::DashMap; -use lasso::{Spur, ThreadedRodeo}; -use time::OffsetDateTime; -use tokio::sync::RwLock; -use tracing::{Level, instrument}; -use uuid::Uuid; - -/// Parse a UUID string into a 16-byte array. -/// Returns None if the string is not a valid UUID. -fn parse_uuid_bytes(s: &str) -> Option<[u8; 16]> { - Uuid::parse_str(s).ok().map(|u| *u.as_bytes()) -} - -/// Format a 16-byte array as a UUID string. -fn format_uuid_bytes(bytes: &[u8; 16]) -> String { - Uuid::from_bytes(*bytes).to_string() -} - -/// Pre-computed frecency data for O(1) lookup. -#[derive(Debug, Clone, Default)] -pub struct FrecencyData { - /// Total number of times this command was used. - pub count: u32, - /// Most recent usage timestamp (unix seconds). - pub last_used: i64, -} - -impl FrecencyData { - /// Record a new usage of this command. - pub fn record_use(&mut self, timestamp: i64) { - self.count += 1; - if timestamp > self.last_used { - self.last_used = timestamp; - } - } - - /// Compute frecency score based on count and recency. - /// - /// Uses a decay function where more recent commands score higher. - /// The formula balances frequency (how often) with recency (how recent). - /// - /// Multipliers allow tuning the relative weights: - /// - `recency_mul`: Multiplier for recency score (default: 1.0) - /// - `frequency_mul`: Multiplier for frequency score (default: 1.0) - /// - /// A multiplier of 0.0 disables that component, 1.0 is unchanged, 2.0 doubles weight. - /// Values like 0.5 reduce weight by half, 1.5 increases by 50%, etc. - #[instrument(level = Level::TRACE, name = "index_frecency_compute")] - pub fn compute(&self, now: i64, recency_mul: f64, frequency_mul: f64) -> u32 { - if self.count == 0 { - return 0; - } - - // Time-based decay: score decreases as time passes - let age_seconds = (now - self.last_used).max(0) as u64; - let age_hours = age_seconds / 3600; - - // Decay factor: recent commands get higher scores - // - Last hour: multiplier ~1.0 - // - Last day: multiplier ~0.5 - // - Last week: multiplier ~0.1 - // - Older: multiplier approaches 0 - let recency_score: f64 = match age_hours { - 0 => 100.0, - 1..=6 => 90.0, - 7..=24 => 70.0, - 25..=72 => 50.0, - 73..=168 => 30.0, - 169..=720 => 15.0, - _ => 5.0, - }; - - // Frequency boost: more uses = higher score (with diminishing returns) - let frequency_score = (f64::from(self.count).ln() * 20.0).min(100.0); - - // Apply multipliers and combine scores, then round to u32 - recency_score - .mul_add(recency_mul, frequency_score * frequency_mul) - .round() as u32 - } -} - -/// Data for a unique command. -pub struct CommandData { - /// History ID of the most recent invocation (16-byte UUID). - most_recent_id: [u8; 16], - /// Timestamp of the most recent invocation. - most_recent_timestamp: i64, - /// Pre-computed global frecency. - pub global_frecency: FrecencyData, - - // Pre-computed indexes for O(1) filter lookups - // Using HashSet instead of DashSet since CommandData lives inside DashMap (already synchronized) - /// All directories where this command has been run (interned keys). - directories: HashSet<Spur>, - /// All hostnames where this command has been run (interned keys). - hosts: HashSet<Spur>, - /// All sessions where this command has been run (as 16-byte UUIDs). - sessions: HashSet<[u8; 16]>, -} - -impl CommandData { - /// Create a new [`CommandData`] from a history entry. - /// Returns None if the history entry has invalid UUIDs. - pub fn new(history: &History, interner: &ThreadedRodeo) -> Option<Self> { - let history_id = parse_uuid_bytes(&history.id.0)?; - let session = parse_uuid_bytes(&history.session)?; - let timestamp = history.timestamp.unix_timestamp(); - - let dir_key = interner.get_or_intern(with_trailing_slash(&history.cwd)); - let host_key = interner.get_or_intern(&history.hostname); - - let mut directories = HashSet::new(); - directories.insert(dir_key); - - let mut hosts = HashSet::new(); - hosts.insert(host_key); - - let mut sessions = HashSet::new(); - sessions.insert(session); - - let mut global_frecency = FrecencyData::default(); - global_frecency.record_use(timestamp); - - Some(Self { - most_recent_id: history_id, - most_recent_timestamp: timestamp, - global_frecency, - directories, - hosts, - sessions, - }) - } - - /// Add an invocation from a history entry. - /// Returns false if the history entry has invalid UUIDs. - pub fn add_invocation(&mut self, history: &History, interner: &ThreadedRodeo) -> bool { - let Some(history_id) = parse_uuid_bytes(&history.id.0) else { - return false; - }; - let Some(session) = parse_uuid_bytes(&history.session) else { - return false; - }; - - let timestamp = history.timestamp.unix_timestamp(); - - // Update global frecency - self.global_frecency.record_use(timestamp); - - // Update pre-computed indexes for O(1) filter lookups - let dir_key = interner.get_or_intern(with_trailing_slash(&history.cwd)); - self.directories.insert(dir_key); - self.hosts.insert(interner.get_or_intern(&history.hostname)); - self.sessions.insert(session); - - // Update most recent if this invocation is newer - if timestamp > self.most_recent_timestamp { - self.most_recent_id = history_id; - self.most_recent_timestamp = timestamp; - } - - true - } - - /// Get the most recent history ID for this command. - pub fn most_recent_id(&self) -> String { - format_uuid_bytes(&self.most_recent_id) - } - - /// Check if any invocation matches a directory filter (exact match). - /// O(1) lookup using pre-computed index. - pub fn has_invocation_in_dir(&self, dir: &str, interner: &ThreadedRodeo) -> bool { - interner - .get(dir) - .is_some_and(|spur| self.directories.contains(&spur)) - } - - /// Check if any invocation matches a directory prefix (workspace/git root). - /// O(n) where n = number of unique directories for this command. - pub fn has_invocation_in_workspace( - &self, - prefix: &str, - interner: &ThreadedRodeo, - ) -> bool { - self.directories - .iter() - .any(|&spur| interner.resolve(&spur).starts_with(prefix)) - } - - /// Check if any invocation matches a hostname. - /// O(1) lookup using pre-computed index. - pub fn has_invocation_on_host(&self, hostname: &str, interner: &ThreadedRodeo) -> bool { - interner - .get(hostname) - .is_some_and(|spur| self.hosts.contains(&spur)) - } - - /// Check if any invocation matches a session. - /// O(1) lookup using pre-computed index. - pub fn has_invocation_in_session(&self, session: &str) -> bool { - parse_uuid_bytes(session).is_some_and(|bytes| self.sessions.contains(&bytes)) - } -} - -/// Filter mode for search queries. -#[derive(Debug, Clone, PartialEq, Eq)] -pub enum IndexFilterMode { - /// No filtering - search all commands. - Global, - /// Filter to commands run in a specific directory. - Directory(String), - /// Filter to commands run in a workspace (directory prefix). - Workspace(String), - /// Filter to commands run on a specific host. - Host(String), - /// Filter to commands run in a specific session. - Session(String), -} - -/// Context for search queries. -#[derive(Debug, Clone, Default)] -pub struct QueryContext { - #[expect(dead_code)] - pub cwd: Option<String>, - #[expect(dead_code)] - pub git_root: Option<String>, - #[expect(dead_code)] - pub hostname: Option<String>, - #[expect(dead_code)] - pub session_id: Option<String>, -} - -/// Shareable frecency map: command -> frecency score. -/// Wrapped in Arc for zero-copy sharing with scorer callbacks. -type FrecencyMap = Arc<HashMap<Arc<str>, u32>>; - -/// A deduplicated search index with frecency-based ranking. -/// -/// Commands are stored by their text, with metadata about all invocations. -/// Nucleo handles fuzzy matching, while frecency is computed via scorer callback. -/// -/// Global frecency is precomputed by a background task and used for scoring. -/// If frecency data is not available, search still works but without frecency ranking; -/// although this should never happen due to precomputing the frecency map. -pub struct SearchIndex { - /// Map from command text to command data. - /// Using `DashMap` for concurrent read/write access, wrapped in Arc for sharing with scorer. - /// Keys are Arc<str> to enable zero-copy sharing with `frecency_map`. - commands: Arc<DashMap<Arc<str>, CommandData>>, - /// Nucleo fuzzy matcher - items are command strings. - nucleo: RwLock<Nucleo<String>>, - /// Injector for adding new commands to Nucleo. - injector: Injector<String>, - /// Precomputed global frecency map. Updated by background task. - frecency_map: RwLock<Option<FrecencyMap>>, - /// String interner for deduplicating cwd, hostname, and directory paths. - interner: Arc<ThreadedRodeo>, -} - -impl SearchIndex { - /// Create a new empty search index. - pub fn new() -> Self { - let nucleo_config = atuin_nucleo::Config::DEFAULT; - // Single column for command text - let nucleo = Nucleo::<String>::new(nucleo_config, Arc::new(|| {}), None, 1); - let injector = nucleo.injector(); - - Self { - commands: Arc::new(DashMap::new()), - nucleo: RwLock::new(nucleo), - injector, - frecency_map: RwLock::new(None), - interner: Arc::new(ThreadedRodeo::new()), - } - } - - /// Add a history entry to the index. - /// - /// If the command already exists, updates its invocation data. - /// If it's a new command, adds it to both the map and Nucleo. - pub fn add_history(&self, history: &History) { - let command = history.command.as_str(); - - // DashMap with Arc<str> keys can be looked up with &str via Borrow trait - if let Some(mut entry) = self.commands.get_mut(command) { - // Existing command - just update invocations - entry.add_invocation(history, &self.interner); - } else { - // New command - create Arc<str> once and share it - let Some(data) = CommandData::new(history, &self.interner) else { - return; // Invalid UUIDs, skip this entry - }; - let command_arc: Arc<str> = command.into(); - self.commands.insert(Arc::clone(&command_arc), data); - // Nucleo still needs String (unavoidable copy for fuzzy matching) - self.injector.push(command_arc.to_string(), |cmd, cols| { - cols[0] = cmd.clone().into(); - }); - } - // Note: frecency_map is rebuilt by background task, not invalidated here - } - - /// Add multiple history entries to the index. - pub fn add_histories(&self, histories: &[History]) { - for history in histories { - self.add_history(history); - } - } - - /// Get the number of unique commands in the index. - pub fn command_count(&self) -> usize { - self.commands.len() - } - - /// Search for commands matching a query. - /// - /// Returns a list of history IDs (most recent invocation per command). - /// Uses precomputed global frecency for scoring if available. - #[instrument(skip_all, level = Level::TRACE, name = "index_search", fields(query = %query))] - #[expect( - clippy::significant_drop_tightening, - reason = "The nucleo early drop is a false-positive" - )] - pub async fn search( - &self, - query: &str, - filter_mode: IndexFilterMode, - // TODO(@bpeetz): Use the query context here <2026-06-12> - #[expect(unused)] context: &QueryContext, - limit: u32, - ) -> Vec<String> { - let mut nucleo = self.nucleo.write().await; - - // Get precomputed frecency map (may be None if not yet computed) - let frecency_map = self.frecency_map.read().await.clone(); - - // Build filter based on mode - let filter = self.build_filter(&filter_mode); - nucleo.set_filter(filter); - - // Build scorer from precomputed frecency (or None if not available) - let scorer = Self::build_scorer(frecency_map); - nucleo.set_scorer(scorer); - - // Update pattern - nucleo.pattern.reparse( - 0, - query, - pattern::CaseMatching::Smart, - pattern::Normalization::Smart, - false, - ); - - tracing::span!(Level::TRACE, "index_search_tick").in_scope(|| { - // Tick until complete - while nucleo.tick(10).running {} - }); - - // Collect results - let snapshot = nucleo.snapshot(); - let matched_count = snapshot.matched_item_count().min(limit); - - tracing::span!(Level::TRACE, "index_search_results").in_scope(|| { - snapshot - .matched_items(..matched_count) - .filter_map(|item| { - let cmd = item.data; - // DashMap<Arc<str>, _>::get accepts &str via Borrow trait - self.commands - .get(cmd.as_str()) - .map(|data| data.most_recent_id()) - }) - .collect() - }) - } - - /// Rebuild the global frecency map. - /// - /// This should be called by a background task periodically. - /// The map is used for scoring search results. - /// - /// Uses multipliers from search settings: - /// - `recency_score_multiplier`: Weight for recency component - /// - `frequency_score_multiplier`: Weight for frequency component - /// - `frecency_score_multiplier`: Overall multiplier for final score - #[instrument(skip_all, level = Level::DEBUG, name = "rebuild_frecency")] - pub async fn rebuild_frecency(&self, search_settings: &Search) { - let now = OffsetDateTime::now_utc().unix_timestamp(); - let mut frecency_map: HashMap<Arc<str>, u32> = HashMap::new(); - - // Clamp multipliers to non-negative values to prevent broken frecency ranking - // (negative values would produce unexpected results when cast to u32) - let recency_mul = search_settings.recency_score_multiplier.max(0.0); - let frequency_mul = search_settings.frequency_score_multiplier.max(0.0); - let frecency_mul = search_settings.frecency_score_multiplier.max(0.0); - - for entry in self.commands.iter() { - let frecency = entry - .global_frecency - .compute(now, recency_mul, frequency_mul); - // Apply overall frecency multiplier and round to u32 - let frecency = (f64::from(frecency) * frecency_mul).round() as u32; - // Arc::clone is cheap - just increments reference count - frecency_map.insert(Arc::clone(entry.key()), frecency); - } - - *self.frecency_map.write().await = Some(Arc::new(frecency_map)); - } - - /// Build filter predicate for the given mode. - fn build_filter(&self, mode: &IndexFilterMode) -> Option<atuin_nucleo::Filter<String>> { - // For Global mode, no filter needed - if matches!(mode, IndexFilterMode::Global) { - return None; - } - - // Pre-compute which commands pass the filter - // Use HashSet<String> for the short-lived filter (simpler than Arc lookup) - let passing_commands: Arc<HashSet<String>> = { - let mut set = HashSet::new(); - for entry in self.commands.iter() { - let passes = match mode { - IndexFilterMode::Global => unreachable!(), - IndexFilterMode::Directory(dir) => { - entry.has_invocation_in_dir(dir, &self.interner) - } - IndexFilterMode::Workspace(prefix) => { - entry.has_invocation_in_workspace(prefix, &self.interner) - } - IndexFilterMode::Host(hostname) => { - entry.has_invocation_on_host(hostname, &self.interner) - } - IndexFilterMode::Session(session) => entry.has_invocation_in_session(session), - }; - if passes { - // Convert Arc<str> to String for filter lookup - set.insert(entry.key().to_string()); - } - } - Arc::new(set) - }; - - Some(Arc::new(move |cmd: &String| passing_commands.contains(cmd))) - } - - /// Build scorer from precomputed frecency map. - /// - /// Returns None if frecency map is not available (search still works, just without frecency ranking). - fn build_scorer(frecency_map: Option<FrecencyMap>) -> Option<atuin_nucleo::Scorer<String>> { - let map = frecency_map?; - Some(Arc::new(move |cmd: &String, fuzzy_score: u32| { - // HashMap<Arc<str>, _>::get accepts &str via Borrow trait - let frecency = map.get(cmd.as_str()).copied().unwrap_or(0); - fuzzy_score + frecency - })) - } -} - -impl Default for SearchIndex { - fn default() -> Self { - Self::new() - } -} - -#[cfg(test)] -mod tests { - use super::FrecencyData; - - #[test] - fn frecency_data_compute() { - let now = 1_000_000i64; - - // Recent command (with default multipliers of 1.0) - let recent = FrecencyData { - count: 5, - last_used: now - 60, // 1 minute ago - }; - assert!(recent.compute(now, 1.0, 1.0) > 100); // High score - - // Old command - let old = FrecencyData { - count: 5, - last_used: now - 86400 * 30, // 30 days ago - }; - assert!(old.compute(now, 1.0, 1.0) < recent.compute(now, 1.0, 1.0)); - - // Frequently used old command - let frequent_old = FrecencyData { - count: 100, - last_used: now - 86400 * 7, // 1 week ago - }; - // Should still have decent score due to frequency - assert!(frequent_old.compute(now, 1.0, 1.0) > 50); - } - - #[test] - fn frecency_data_compute_with_multipliers() { - let now = 1_000_000_i64; - - let data = FrecencyData { - count: 5, - last_used: now - 60, // 1 minute ago (recency_score = 100) - }; - - // Default multipliers (1.0, 1.0) - let default_score = data.compute(now, 1.0, 1.0); - - // Double recency weight - let double_recency = data.compute(now, 2.0, 1.0); - assert!(double_recency > default_score); - - // Double frequency weight - let double_frequency = data.compute(now, 1.0, 2.0); - assert!(double_frequency > default_score); - - // Zero out recency (only frequency counts) - let no_recency = data.compute(now, 0.0, 1.0); - assert!(no_recency < default_score); - - // Zero out frequency (only recency counts) - let no_frequency = data.compute(now, 1.0, 0.0); - assert!(no_frequency < default_score); - - // Zero both (should be zero) - let no_score = data.compute(now, 0.0, 0.0); - assert_eq!(no_score, 0); - - // Fractional multipliers - let half_recency = data.compute(now, 0.5, 1.0); - assert!(half_recency < default_score); - assert!(half_recency > no_recency); - - // 1.5x multiplier - let boost_recency = data.compute(now, 1.5, 1.0); - assert!(boost_recency > default_score); - assert!(boost_recency < double_recency); - } -} diff --git a/crates/daemon/src/server.rs b/crates/daemon/src/server.rs index 335f8260..97c4fe48 100644 --- a/crates/daemon/src/server.rs +++ b/crates/daemon/src/server.rs @@ -1,18 +1,13 @@ use eyre::Result; use crate::{ - atuin_client::settings::Settings, - atuin_daemon::{ - components::{ - history::HistoryGrpcService, search::SearchGrpcService, semantic::SemanticGrpcService, - }, - daemon::DaemonHandle, - generated::{ - control::{ControlService, control_server::ControlServer}, - history::history_server::HistoryServer, - search::search_server::SearchServer, - semantic::semantic_server::SemanticServer, - }, + aclient::settings::Settings, + components::{history::HistoryGrpcService, semantic::SemanticGrpcService}, + daemon::DaemonHandle, + generated::{ + control::{ControlService, control_server::ControlServer}, + history::history_server::HistoryServer, + semantic::semantic_server::SemanticServer, }, }; @@ -21,10 +16,9 @@ use crate::{ /// This starts the gRPC server in the background and returns immediately. /// The server will shut down when a [`ShutdownRequested`] event is received. #[cfg(unix)] -pub fn run_grpc_server( +pub(crate) fn run_grpc_server( settings: &Settings, history_service: HistoryServer<HistoryGrpcService>, - search_service: SearchServer<SearchGrpcService>, semantic_service: SemanticServer<SemanticGrpcService>, control_service: ControlServer<ControlService>, handle: DaemonHandle, @@ -107,7 +101,6 @@ pub fn run_grpc_server( if let Err(e) = Server::builder() .add_service(history_service) - .add_service(search_service) .add_service(semantic_service) .add_service(control_service) .serve_with_incoming_shutdown(uds_stream, shutdown_signal) |
