diff options
Diffstat (limited to 'crates/server')
33 files changed, 1252 insertions, 0 deletions
diff --git a/crates/server/Cargo.toml b/crates/server/Cargo.toml new file mode 100644 index 00000000..982f32b1 --- /dev/null +++ b/crates/server/Cargo.toml @@ -0,0 +1,103 @@ +[package] +name = "turtle-server" +edition = "2024" +description = "sync server for turtle history db" +readme = "./README.md" + +rust-version = { workspace = true } +version = { workspace = true } +authors = { workspace = true } +license = { workspace = true } +homepage = { workspace = true } +repository = { workspace = true } + +[dependencies] +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"] } +clap_complete = "4.5.8" +clap_complete_nushell = "4.5.4" +colored = "2.0.4" +config = { version = "0.15.8", default-features = false, features = ["toml"] } +crossterm = {version = "0.29.0", features = ["use-dev-tty", "serde"] } +crypto_secretbox = "0.1.1" +dashmap = "6.1.0" +directories = "6.0.0" +eyre = "0.6" +fs-err = "3.1" +fs4 = "0.13.1" +fuzzy-matcher = "0.3.7" +hyper-util = "0.1" +indicatif = "0.18.0" +interim = { version = "0.2.0", features = ["time_0_3"] } +itertools = "0.14.0" +lasso = { version = "0.7", features = ["multi-threaded"] } +log = "0.4" +metrics = "0.24" +metrics-exporter-prometheus = { version = "0.18", default-features = false } +minspan = "0.1.5" +norm = { version = "0.1.1", features = ["fzf-v2"] } +notify = "7" +prost = "0.14" +rand = { version = "0.8.5", features = ["std"] } +ratatui = "0.30.0" +regex = "1.10.5" +reqwest = { version = "0.13", features = ["json", "rustls-no-provider", "stream"], default-features = false } +rmp = { version = "0.8.14" } +runtime-format = "0.1.3" +rustix = { version = "1.1.4", features = ["process", "fs"] } +rustls = { version = "0.23", default-features = false, features = [ "ring", "std", "tls12", ] } +rusty_paserk = { version = "0.5.0", default-features = false, features = [ "v4", "serde", ] } +rusty_paseto = { version = "0.8.0", default-features = false } +semver = "1.0.20" +serde = { version = "1.0.202", features = ["derive"] } +serde_json = "1.0.119" +serde_regex = "1.1.0" +serde_with = "3.8.1" +shellexpand = "3" +sql-builder = "3" +sqlx = { version = "0.8", features = ["runtime-tokio-rustls", "time", "postgres", "uuid", "sqlite", "regexp"] } +thiserror = "2" +time = { version = "0.3.47", features = [ "serde-human-readable", "macros", "local-offset", "macros", "formatting", "parsing"] } +tokio = { version = "1", features = ["full"] } +tokio-stream = { version = "0.1.14", features = ["net"] } +toml_edit = "0.25.4" +tonic = "0.14" +tonic-prost = "0.14" +tower = "0.5" +tower-http = { version = "0.6", features = ["trace"] } +tracing = "0.1" +tracing-appender = "0.2" +tracing-subscriber = { version = "0.3", features = ["ansi", "fmt", "registry", "env-filter", "json"] } +typed-builder = "0.18.2" +unicode-segmentation = "1.11.0" +unicode-width = "0.2" +url = "2.5.2" +uuid = { version = "1.9", features = ["v4", "v7", "serde"] } +vt100 = "0.16" +whoami = "2.1.0" + +[target.'cfg(target_os = "linux")'.dependencies] +arboard = { version = "3.4", default-features = false, features = [ "wayland-data-control", ] } +listenfd = "1.0.1" + +[target.'cfg(unix)'.dependencies] +daemonize = "0.5.0" +portable-pty = "0.9" +signal-hook = "0.3" + +[dev-dependencies] +tokio = { version = "1", features = ["full"] } + +[build-dependencies] +protox = "0.9" +tonic-prost-build = "0.14" + +[package.metadata.docs.rs] +all-features = true + +[lints] +workspace = true diff --git a/crates/server/db/server-pg-migrations/20210425153745_create_history.sql b/crates/server/db/server-pg-migrations/20210425153745_create_history.sql new file mode 100644 index 00000000..2c2d17b0 --- /dev/null +++ b/crates/server/db/server-pg-migrations/20210425153745_create_history.sql @@ -0,0 +1,11 @@ +create table history ( + id bigserial primary key, + client_id text not null unique, -- the client-generated ID + user_id bigserial not null, -- allow multiple users + hostname text not null, -- a unique identifier from the client (can be hashed, random, whatever) + timestamp timestamp not null, -- one of the few non-encrypted metadatas + + data varchar(8192) not null, -- store the actual history data, encrypted. I don't wanna know! + + created_at timestamp not null default current_timestamp +); diff --git a/crates/server/db/server-pg-migrations/20210425153757_create_users.sql b/crates/server/db/server-pg-migrations/20210425153757_create_users.sql new file mode 100644 index 00000000..a25dcced --- /dev/null +++ b/crates/server/db/server-pg-migrations/20210425153757_create_users.sql @@ -0,0 +1,10 @@ +create table users ( + id bigserial primary key, -- also store our own ID + username varchar(32) not null unique, -- being able to contact users is useful + email varchar(128) not null unique, -- being able to contact users is useful + password varchar(128) not null unique +); + +-- the prior index is case sensitive :( +CREATE UNIQUE INDEX email_unique_idx on users (LOWER(email)); +CREATE UNIQUE INDEX username_unique_idx on users (LOWER(username)); diff --git a/crates/server/db/server-pg-migrations/20210425153800_create_sessions.sql b/crates/server/db/server-pg-migrations/20210425153800_create_sessions.sql new file mode 100644 index 00000000..c2fb6559 --- /dev/null +++ b/crates/server/db/server-pg-migrations/20210425153800_create_sessions.sql @@ -0,0 +1,6 @@ +-- Add migration script here +create table sessions ( + id bigserial primary key, + user_id bigserial, + token varchar(128) unique not null +); diff --git a/crates/server/db/server-pg-migrations/20220419082412_add_count_trigger.sql b/crates/server/db/server-pg-migrations/20220419082412_add_count_trigger.sql new file mode 100644 index 00000000..dd1afa88 --- /dev/null +++ b/crates/server/db/server-pg-migrations/20220419082412_add_count_trigger.sql @@ -0,0 +1,51 @@ +-- Prior to this, the count endpoint was super naive and just ran COUNT(1). +-- This is slow asf. Now that we have an amount of actual traffic, +-- stop doing that! +-- This basically maintains a count, so we can read ONE row, instead of ALL the +-- rows. Much better. +-- Future optimisation could use some sort of cache so we don't even need to hit +-- postgres at all. + +create table total_history_count_user( + id bigserial primary key, + user_id bigserial, + total integer -- try and avoid using keywords - hence total, not count +); + +create or replace function user_history_count() +returns trigger as +$func$ +begin + if (TG_OP='INSERT') then + update total_history_count_user set total = total + 1 where user_id = new.user_id; + + if not found then + insert into total_history_count_user(user_id, total) + values ( + new.user_id, + (select count(1) from history where user_id = new.user_id) + ); + end if; + + elsif (TG_OP='DELETE') then + update total_history_count_user set total = total - 1 where user_id = new.user_id; + + if not found then + insert into total_history_count_user(user_id, total) + values ( + new.user_id, + (select count(1) from history where user_id = new.user_id) + ); + end if; + end if; + + return NEW; -- this is actually ignored for an after trigger, but oh well +end; +$func$ +language plpgsql volatile -- pldfplplpflh +cost 100; -- default value + +create trigger tg_user_history_count + after insert or delete on history + for each row + execute procedure user_history_count(); diff --git a/crates/server/db/server-pg-migrations/20220421073605_fix_count_trigger_delete.sql b/crates/server/db/server-pg-migrations/20220421073605_fix_count_trigger_delete.sql new file mode 100644 index 00000000..6198f300 --- /dev/null +++ b/crates/server/db/server-pg-migrations/20220421073605_fix_count_trigger_delete.sql @@ -0,0 +1,35 @@ +-- the old version of this function used NEW in the delete part when it should +-- use OLD + +create or replace function user_history_count() +returns trigger as +$func$ +begin + if (TG_OP='INSERT') then + update total_history_count_user set total = total + 1 where user_id = new.user_id; + + if not found then + insert into total_history_count_user(user_id, total) + values ( + new.user_id, + (select count(1) from history where user_id = new.user_id) + ); + end if; + + elsif (TG_OP='DELETE') then + update total_history_count_user set total = total - 1 where user_id = old.user_id; + + if not found then + insert into total_history_count_user(user_id, total) + values ( + old.user_id, + (select count(1) from history where user_id = old.user_id) + ); + end if; + end if; + + return NEW; -- this is actually ignored for an after trigger, but oh well +end; +$func$ +language plpgsql volatile -- pldfplplpflh +cost 100; -- default value diff --git a/crates/server/db/server-pg-migrations/20220421174016_larger-commands.sql b/crates/server/db/server-pg-migrations/20220421174016_larger-commands.sql new file mode 100644 index 00000000..0ac43433 --- /dev/null +++ b/crates/server/db/server-pg-migrations/20220421174016_larger-commands.sql @@ -0,0 +1,3 @@ +-- Make it 4x larger. Most commands are less than this, but as it's base64 +-- SOME are more than 8192. Should be enough for now. +ALTER TABLE history ALTER COLUMN data TYPE varchar(32768); diff --git a/crates/server/db/server-pg-migrations/20220426172813_user-created-at.sql b/crates/server/db/server-pg-migrations/20220426172813_user-created-at.sql new file mode 100644 index 00000000..a9138194 --- /dev/null +++ b/crates/server/db/server-pg-migrations/20220426172813_user-created-at.sql @@ -0,0 +1 @@ +alter table users add column created_at timestamp not null default now(); diff --git a/crates/server/db/server-pg-migrations/20220505082442_create-events.sql b/crates/server/db/server-pg-migrations/20220505082442_create-events.sql new file mode 100644 index 00000000..57e16ec7 --- /dev/null +++ b/crates/server/db/server-pg-migrations/20220505082442_create-events.sql @@ -0,0 +1,14 @@ +create type event_type as enum ('create', 'delete'); + +create table events ( + id bigserial primary key, + client_id text not null unique, -- the client-generated ID + user_id bigserial not null, -- allow multiple users + hostname text not null, -- a unique identifier from the client (can be hashed, random, whatever) + timestamp timestamp not null, -- one of the few non-encrypted metadatas + + event_type event_type, + data text not null, -- store the actual history data, encrypted. I don't wanna know! + + created_at timestamp not null default current_timestamp +); diff --git a/crates/server/db/server-pg-migrations/20220610074049_history-length.sql b/crates/server/db/server-pg-migrations/20220610074049_history-length.sql new file mode 100644 index 00000000..b1c23016 --- /dev/null +++ b/crates/server/db/server-pg-migrations/20220610074049_history-length.sql @@ -0,0 +1,2 @@ +-- Add migration script here +alter table history alter column data type text; diff --git a/crates/server/db/server-pg-migrations/20230315220537_drop-events.sql b/crates/server/db/server-pg-migrations/20230315220537_drop-events.sql new file mode 100644 index 00000000..fe3cae17 --- /dev/null +++ b/crates/server/db/server-pg-migrations/20230315220537_drop-events.sql @@ -0,0 +1,2 @@ +-- Add migration script here +drop table events; diff --git a/crates/server/db/server-pg-migrations/20230315224203_create-deleted.sql b/crates/server/db/server-pg-migrations/20230315224203_create-deleted.sql new file mode 100644 index 00000000..9a9e6263 --- /dev/null +++ b/crates/server/db/server-pg-migrations/20230315224203_create-deleted.sql @@ -0,0 +1,5 @@ +-- Add migration script here +alter table history add column if not exists deleted_at timestamp; + +-- queries will all be selecting the ids of history for a user, that has been deleted +create index if not exists history_deleted_index on history(client_id, user_id, deleted_at); diff --git a/crates/server/db/server-pg-migrations/20230515221038_trigger-delete-only.sql b/crates/server/db/server-pg-migrations/20230515221038_trigger-delete-only.sql new file mode 100644 index 00000000..3d0bba52 --- /dev/null +++ b/crates/server/db/server-pg-migrations/20230515221038_trigger-delete-only.sql @@ -0,0 +1,30 @@ +-- We do not need to run the trigger on deletes, as the only time we are deleting history is when the user +-- has already been deleted +-- This actually slows down deleting all the history a good bit! + +create or replace function user_history_count() +returns trigger as +$func$ +begin + if (TG_OP='INSERT') then + update total_history_count_user set total = total + 1 where user_id = new.user_id; + + if not found then + insert into total_history_count_user(user_id, total) + values ( + new.user_id, + (select count(1) from history where user_id = new.user_id) + ); + end if; + end if; + + return NEW; -- this is actually ignored for an after trigger, but oh well +end; +$func$ +language plpgsql volatile -- pldfplplpflh +cost 100; -- default value + +create or replace trigger tg_user_history_count + after insert on history + for each row + execute procedure user_history_count(); diff --git a/crates/server/db/server-pg-migrations/20230623070418_records.sql b/crates/server/db/server-pg-migrations/20230623070418_records.sql new file mode 100644 index 00000000..a3e5de2e --- /dev/null +++ b/crates/server/db/server-pg-migrations/20230623070418_records.sql @@ -0,0 +1,15 @@ +-- Add migration script here +create table records ( + id uuid primary key, -- remember to use uuidv7 for happy indices <3 + client_id uuid not null, -- I am too uncomfortable with the idea of a client-generated primary key + host uuid not null, -- a unique identifier for the host + parent uuid default null, -- the ID of the parent record, bearing in mind this is a linked list + timestamp bigint not null, -- not a timestamp type, as those do not have nanosecond precision + version text not null, + tag text not null, -- what is this? history, kv, whatever. Remember clients get a log per tag per host + data text not null, -- store the actual history data, encrypted. I don't wanna know! + cek text not null, + + user_id bigint not null, -- allow multiple users + created_at timestamp not null default current_timestamp +); diff --git a/crates/server/db/server-pg-migrations/20231202170508_create-store.sql b/crates/server/db/server-pg-migrations/20231202170508_create-store.sql new file mode 100644 index 00000000..ffb57966 --- /dev/null +++ b/crates/server/db/server-pg-migrations/20231202170508_create-store.sql @@ -0,0 +1,15 @@ +-- Add migration script here +create table store ( + id uuid primary key, -- remember to use uuidv7 for happy indices <3 + client_id uuid not null, -- I am too uncomfortable with the idea of a client-generated primary key, even though it's fine mathematically + host uuid not null, -- a unique identifier for the host + idx bigint not null, -- the index of the record in this store, identified by (host, tag) + timestamp bigint not null, -- not a timestamp type, as those do not have nanosecond precision + version text not null, + tag text not null, -- what is this? history, kv, whatever. Remember clients get a log per tag per host + data text not null, -- store the actual history data, encrypted. I don't wanna know! + cek text not null, + + user_id bigint not null, -- allow multiple users + created_at timestamp not null default current_timestamp +); diff --git a/crates/server/db/server-pg-migrations/20231203124112_create-store-idx.sql b/crates/server/db/server-pg-migrations/20231203124112_create-store-idx.sql new file mode 100644 index 00000000..56d67145 --- /dev/null +++ b/crates/server/db/server-pg-migrations/20231203124112_create-store-idx.sql @@ -0,0 +1,2 @@ +-- Add migration script here +create unique index record_uniq ON store(user_id, host, tag, idx); diff --git a/crates/server/db/server-pg-migrations/20240108124837_drop-some-defaults.sql b/crates/server/db/server-pg-migrations/20240108124837_drop-some-defaults.sql new file mode 100644 index 00000000..ad2af5a1 --- /dev/null +++ b/crates/server/db/server-pg-migrations/20240108124837_drop-some-defaults.sql @@ -0,0 +1,4 @@ +-- Add migration script here +alter table history alter column user_id drop default; +alter table sessions alter column user_id drop default; +alter table total_history_count_user alter column user_id drop default; diff --git a/crates/server/db/server-pg-migrations/20240614104159_idx-cache.sql b/crates/server/db/server-pg-migrations/20240614104159_idx-cache.sql new file mode 100644 index 00000000..12bbbecc --- /dev/null +++ b/crates/server/db/server-pg-migrations/20240614104159_idx-cache.sql @@ -0,0 +1,8 @@ +CREATE TABLE store_idx_cache( + id BIGSERIAL PRIMARY KEY, + user_id UUID, + + host UUID, + tag TEXT, + idx BIGINT +); diff --git a/crates/server/db/server-pg-migrations/20240621110731_user-verified.sql b/crates/server/db/server-pg-migrations/20240621110731_user-verified.sql new file mode 100644 index 00000000..6eba02ec --- /dev/null +++ b/crates/server/db/server-pg-migrations/20240621110731_user-verified.sql @@ -0,0 +1,8 @@ +alter table users add verified_at timestamp with time zone default null; + +create table user_verification_token( + id bigserial primary key, + user_id bigint unique references users(id), + token text, + valid_until timestamp with time zone +); diff --git a/crates/server/db/server-pg-migrations/20240702094825_idx_cache_index.sql b/crates/server/db/server-pg-migrations/20240702094825_idx_cache_index.sql new file mode 100644 index 00000000..d1a7b194 --- /dev/null +++ b/crates/server/db/server-pg-migrations/20240702094825_idx_cache_index.sql @@ -0,0 +1 @@ +create unique index store_idx_cache_uniq on store_idx_cache(user_id, host, tag); diff --git a/crates/server/db/server-pg-migrations/20260127000000_remove-email-verification.sql b/crates/server/db/server-pg-migrations/20260127000000_remove-email-verification.sql new file mode 100644 index 00000000..15309920 --- /dev/null +++ b/crates/server/db/server-pg-migrations/20260127000000_remove-email-verification.sql @@ -0,0 +1,2 @@ +drop table if exists user_verification_token; +alter table users drop column if exists verified_at; diff --git a/crates/server/db/server-pg-migrations/20260611222503_make_user-id_an_uuid.sql b/crates/server/db/server-pg-migrations/20260611222503_make_user-id_an_uuid.sql new file mode 100644 index 00000000..d31c23e2 --- /dev/null +++ b/crates/server/db/server-pg-migrations/20260611222503_make_user-id_an_uuid.sql @@ -0,0 +1,11 @@ +-- Add migration script here + +ALTER TABLE records +DROP COLUMN IF EXISTS user_id; +ALTER TABLE records +ADD COLUMN user_id UUID NOT NULL; + +ALTER TABLE store +DROP COLUMN IF EXISTS user_id; +ALTER TABLE store +ADD COLUMN user_id UUID NOT NULL; diff --git a/crates/server/src/database/db/mod.rs b/crates/server/src/database/db/mod.rs new file mode 100644 index 00000000..77bd0c61 --- /dev/null +++ b/crates/server/src/database/db/mod.rs @@ -0,0 +1,274 @@ +use std::collections::HashMap; + +use rand::Rng; + +use crate::{ + atuin_common::record::{EncryptedData, HostId, Record, RecordIdx, RecordStatus}, + atuin_server::database::{DbError, DbResult, DbSettings, models::User}, +}; +use sqlx::postgres::PgPoolOptions; + +use tracing::instrument; +use uuid::Uuid; +use wrappers::DbRecord; + +mod wrappers; + +const MIN_PG_VERSION: u32 = 14; + +#[derive(Clone)] +pub(crate) struct ServerPostgres { + pool: sqlx::Pool<sqlx::postgres::Postgres>, + /// Optional read replica pool for read-only queries + read_pool: Option<sqlx::Pool<sqlx::postgres::Postgres>>, +} + +impl ServerPostgres { + /// Returns the appropriate pool for read operations. + /// Uses `read_pool` if available, otherwise falls back to the primary pool. + fn read_pool(&self) -> &sqlx::Pool<sqlx::postgres::Postgres> { + self.read_pool.as_ref().unwrap_or(&self.pool) + } +} + +impl ServerPostgres { + pub(crate) async fn new(settings: &DbSettings) -> DbResult<Self> { + let pool = PgPoolOptions::new() + .max_connections(100) + .connect(settings.db_uri.as_str()) + .await?; + + // Call server_version_num to get the DB server's major version number + // The call returns None for servers older than 8.x. + let pg_major_version: u32 = + pool.acquire() + .await? + .server_version_num() + .ok_or(DbError::Other(eyre::Report::msg( + "could not get PostgreSQL version", + )))? + / 10000; + + if pg_major_version < MIN_PG_VERSION { + return Err(DbError::Other(eyre::Report::msg(format!( + "unsupported PostgreSQL version {pg_major_version}, minimum required is {MIN_PG_VERSION}" + )))); + } + + sqlx::migrate!("./db/server-pg-migrations") + .run(&pool) + .await + .map_err(|error| DbError::Other(error.into()))?; + + // Create read replica pool if configured + let read_pool = if let Some(read_db_uri) = &settings.read_db_uri { + tracing::info!("Connecting to read replica database"); + let read_pool = PgPoolOptions::new() + .max_connections(100) + .connect(read_db_uri.as_str()) + .await?; + + // Verify the read replica is also a supported PostgreSQL version + let read_pg_major_version: u32 = read_pool + .acquire() + .await? + .server_version_num() + .ok_or(DbError::Other(eyre::Report::msg( + "could not get PostgreSQL version from read replica", + )))? + / 10000; + + if read_pg_major_version < MIN_PG_VERSION { + return Err(DbError::Other(eyre::Report::msg(format!( + "unsupported PostgreSQL version {read_pg_major_version} on read replica, minimum required is {MIN_PG_VERSION}" + )))); + } + + Some(read_pool) + } else { + None + }; + + Ok(Self { pool, read_pool }) + } + + #[instrument(skip_all)] + pub(crate) async fn add_records( + &self, + user: &User, + records: &[Record<EncryptedData>], + ) -> DbResult<()> { + let mut tx = self.pool.begin().await?; + + // We won't have uploaded this data if it wasn't the max. Therefore, we can deduce the max + // idx without having to make further database queries. Doing the query on this small + // amount of data should be much, much faster. + // + // Worst case, say we get this wrong. We end up caching data that isn't actually the max + // idx, so clients upload again. The cache logic can be verified with a sql query anyway :) + + let mut heads = HashMap::<(HostId, &str), u64>::new(); + + for i in records { + let id = crate::atuin_common::utils::uuid_v7(); + + let result = sqlx::query( + " + INSERT INTO store (id, client_id, host, idx, timestamp, version, tag, data, cek, user_id) + VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10) + ON conflict DO nothing + ", + ) + .bind(id) + .bind(i.id) + .bind(i.host.id) + .bind(i.idx as i64) + .bind(i.timestamp as i64) // throwing away some data, but i64 is still big in terms of time + .bind(&i.version) + .bind(&i.tag) + .bind(&i.data.data) + .bind(&i.data.content_encryption_key) + .bind(user.id) + .execute(&mut *tx) + .await?; + + // Only update heads if we actually inserted the record + if result.rows_affected() > 0 { + heads + .entry((i.host.id, &i.tag)) + .and_modify(|e| { + if i.idx > *e { + *e = i.idx; + } + }) + .or_insert(i.idx); + } + } + + // we've built the map of heads for this push, so commit it to the database + for ((host, tag), idx) in heads { + sqlx::query( + " + INSERT INTO store_idx_cache (user_id, host, tag, idx) + VALUES ($1, $2, $3, $4) + ON conflict(user_id, host, tag) DO update + SET idx = greatest(store_idx_cache.idx, $4) + ", + ) + .bind(user.id) + .bind(host) + .bind(tag) + .bind(idx as i64) + .execute(&mut *tx) + .await?; + } + + tx.commit().await?; + + Ok(()) + } + + #[instrument(skip_all)] + pub(crate) async fn next_records( + &self, + user: &User, + host: HostId, + tag: String, + start: Option<RecordIdx>, + count: u64, + ) -> DbResult<Vec<Record<EncryptedData>>> { + tracing::debug!("{:?} - {:?} - {:?}", host, tag, start); + let start = start.unwrap_or(0); + + let records: Result<Vec<DbRecord>, DbError> = sqlx::query_as( + " + SELECT client_id, host, idx, timestamp, version, tag, data, cek FROM store + WHERE user_id = $1 + AND tag = $2 + AND host = $3 + AND idx >= $4 + ORDER BY idx asc + LIMIT $5 + ", + ) + .bind(user.id) + .bind(tag.clone()) + .bind(host) + .bind(start as i64) + .bind(count as i64) + .fetch_all(self.read_pool()) + .await + .map_err(Into::into); + + let ret = match records { + Ok(records) => { + let records: Vec<Record<EncryptedData>> = records + .into_iter() + .map(|f| { + let record: Record<EncryptedData> = f.into(); + record + }) + .collect(); + + records + } + Err(DbError::NotFound) => { + tracing::debug!("no records found in store: {:?}/{}", host, tag); + return Ok(vec![]); + } + Err(e) => return Err(e), + }; + + Ok(ret) + } + + pub(crate) async fn status(&self, user: &User) -> DbResult<RecordStatus> { + // If IDX_CACHE_ROLLOUT is set, then we + // 1. Read the value of the var, use it as a % chance of using the cache + // 2. If we use the cache, just read from the cache table + // 3. If we don't use the cache, read from the store table + // IDX_CACHE_ROLLOUT should be between 0 and 100. + + let idx_cache_rollout = + std::env::var("IDX_CACHE_ROLLOUT").unwrap_or_else(|_| "0".to_string()); + let idx_cache_rollout = idx_cache_rollout.parse::<f64>().unwrap_or(0.0); + let use_idx_cache = rand::thread_rng().gen_bool(idx_cache_rollout / 100.0); + + let mut res: Vec<(Uuid, String, i64)> = if use_idx_cache { + tracing::debug!("using idx cache for user {}", user.id); + sqlx::query_as( + " + SELECT host, tag, idx + FROM store_idx_cache + WHERE user_id = $1 + ", + ) + .bind(user.id) + .fetch_all(self.read_pool()) + .await? + } else { + tracing::debug!("using aggregate query for user {}", user.id); + sqlx::query_as( + " + SELECT host, tag, max(idx) + FROM store + WHERE user_id = $1 + GROUP BY host, tag + ", + ) + .bind(user.id) + .fetch_all(self.read_pool()) + .await? + }; + + res.sort(); + + let mut status = RecordStatus::new(); + + for i in &res { + status.set_raw(HostId(i.0), i.1.clone(), i.2 as u64); + } + + Ok(status) + } +} diff --git a/crates/server/src/database/db/wrappers.rs b/crates/server/src/database/db/wrappers.rs new file mode 100644 index 00000000..0315e331 --- /dev/null +++ b/crates/server/src/database/db/wrappers.rs @@ -0,0 +1,32 @@ +use crate::atuin_common::record::{EncryptedData, Host, Record}; +use sqlx::{Row, postgres::PgRow}; + +pub(crate) struct DbRecord(pub Record<EncryptedData>); + +impl<'a> ::sqlx::FromRow<'a, PgRow> for DbRecord { + fn from_row(row: &'a PgRow) -> ::sqlx::Result<Self> { + let timestamp: i64 = row.try_get("timestamp")?; + let idx: i64 = row.try_get("idx")?; + + let data = EncryptedData { + data: row.try_get("data")?, + content_encryption_key: row.try_get("cek")?, + }; + + Ok(Self(Record { + id: row.try_get("client_id")?, + host: Host::new(row.try_get("host")?), + idx: idx as u64, + timestamp: timestamp as u64, + version: row.try_get("version")?, + tag: row.try_get("tag")?, + data, + })) + } +} + +impl From<DbRecord> for Record<EncryptedData> { + fn from(other: DbRecord) -> Self { + other.0 + } +} diff --git a/crates/server/src/database/mod.rs b/crates/server/src/database/mod.rs new file mode 100644 index 00000000..43fe5c3b --- /dev/null +++ b/crates/server/src/database/mod.rs @@ -0,0 +1,99 @@ +pub(crate) mod db; +pub(crate) mod models; + +use std::fmt::{Debug, Display}; + +use serde::{Deserialize, Serialize}; + +#[derive(Debug)] +pub(crate) enum DbError { + NotFound, + Other(eyre::Report), +} + +impl Display for DbError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::NotFound => write!(f, "Not found"), + Self::Other(report) => write!(f, "Other: {report}"), + } + } +} + +impl From<time::error::ComponentRange> for DbError { + fn from(error: time::error::ComponentRange) -> Self { + Self::Other(error.into()) + } +} + +impl From<time::error::Error> for DbError { + fn from(error: time::error::Error) -> Self { + Self::Other(error.into()) + } +} + +impl From<sqlx::Error> for DbError { + fn from(error: sqlx::Error) -> Self { + match error { + sqlx::Error::RowNotFound => Self::NotFound, + error => Self::Other(error.into()), + } + } +} + +impl std::error::Error for DbError {} + +pub(crate) type DbResult<T> = Result<T, DbError>; + +#[derive(Debug, PartialEq)] +pub(crate) enum DbType { + Postgres, + Unknown, +} + +#[derive(Clone, Deserialize, Serialize)] +pub(crate) struct DbSettings { + pub(crate) db_uri: String, + + /// Optional URI for read replicas. If set, read-only queries will use this connection. + pub(crate) read_db_uri: Option<String>, +} + +impl DbSettings { + pub(crate) fn db_type(&self) -> DbType { + if self.db_uri.starts_with("postgres://") || self.db_uri.starts_with("postgresql://") { + DbType::Postgres + } else { + DbType::Unknown + } + } +} + +fn redact_db_uri(uri: &str) -> String { + url::Url::parse(uri).map_or_else( + |_| uri.to_string(), + |mut url| { + url.set_password(Some("****")).expect("should be possible"); + url.to_string() + }, + ) +} + +// Do our best to redact passwords so they're not logged in the event of an error. +impl Debug for DbSettings { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + if self.db_type() == DbType::Postgres { + let redacted_uri = redact_db_uri(&self.db_uri); + let redacted_read_uri = self.read_db_uri.as_ref().map(|uri| redact_db_uri(uri)); + f.debug_struct("DbSettings") + .field("db_uri", &redacted_uri) + .field("read_db_uri", &redacted_read_uri) + .finish() + } else { + f.debug_struct("DbSettings") + .field("db_uri", &self.db_uri) + .field("read_db_uri", &self.read_db_uri) + .finish() + } + } +} diff --git a/crates/server/src/database/models.rs b/crates/server/src/database/models.rs new file mode 100644 index 00000000..3fa6f471 --- /dev/null +++ b/crates/server/src/database/models.rs @@ -0,0 +1,5 @@ +use uuid::Uuid; + +pub(crate) struct User { + pub(crate) id: Uuid, +} diff --git a/crates/server/src/handlers/mod.rs b/crates/server/src/handlers/mod.rs new file mode 100644 index 00000000..c4332f80 --- /dev/null +++ b/crates/server/src/handlers/mod.rs @@ -0,0 +1,54 @@ +use crate::atuin_common::api::{ErrorResponse, IndexResponse}; +use axum::{Json, extract::State, http, response::IntoResponse}; + +use crate::router::AppState; + +pub(crate) mod v0; + +const VERSION: &str = env!("CARGO_PKG_VERSION"); + +pub(crate) async fn index(state: State<AppState>) -> Json<IndexResponse> { + let homage = r#""Through the fathomless deeps of space swims the star turtle Great A'Tuin, bearing on its back the four giant elephants who carry on their shoulders the mass of the Discworld." -- Sir Terry Pratchett"#; + + let version = state + .settings + .fake_version + .clone() + .unwrap_or_else(|| VERSION.to_string()); + + Json(IndexResponse { + homage: homage.to_string(), + version, + }) +} + +impl IntoResponse for ErrorResponseStatus<'_> { + fn into_response(self) -> axum::response::Response { + (self.status, Json(self.error)).into_response() + } +} + +pub(crate) struct ErrorResponseStatus<'a> { + pub(crate) error: ErrorResponse<'a>, + pub(crate) status: http::StatusCode, +} + +pub(crate) trait RespExt<'a> { + fn with_status(self, status: http::StatusCode) -> ErrorResponseStatus<'a>; + fn reply(reason: &'a str) -> Self; +} + +impl<'a> RespExt<'a> for ErrorResponse<'a> { + fn with_status(self, status: http::StatusCode) -> ErrorResponseStatus<'a> { + ErrorResponseStatus { + error: self, + status, + } + } + + fn reply(reason: &'a str) -> Self { + Self { + reason: reason.into(), + } + } +} diff --git a/crates/server/src/handlers/v0/mod.rs b/crates/server/src/handlers/v0/mod.rs new file mode 100644 index 00000000..78fb47b8 --- /dev/null +++ b/crates/server/src/handlers/v0/mod.rs @@ -0,0 +1 @@ +pub(crate) mod record; diff --git a/crates/server/src/handlers/v0/record.rs b/crates/server/src/handlers/v0/record.rs new file mode 100644 index 00000000..0381ded8 --- /dev/null +++ b/crates/server/src/handlers/v0/record.rs @@ -0,0 +1,113 @@ +use axum::{Json, extract::Query, extract::State, http::StatusCode}; +use metrics::counter; +use serde::Deserialize; +use tracing::{error, instrument}; + +use crate::{ + handlers::{ErrorResponse, ErrorResponseStatus, RespExt}, + router::{AppState, UserAuth}, +}; + +use crate::atuin_common::record::{EncryptedData, HostId, Record, RecordIdx, RecordStatus}; + +#[instrument(skip_all, fields(user.id = user.id.to_string()))] +pub(crate) async fn post( + UserAuth(user): UserAuth, + state: State<AppState>, + Json(records): Json<Vec<Record<EncryptedData>>>, +) -> Result<(), ErrorResponseStatus<'static>> { + let State(AppState { database, settings }) = state; + + tracing::debug!( + count = records.len(), + user = user.id.to_string(), + "request to add records" + ); + + counter!("atuin_record_uploaded").increment(records.len() as u64); + + let keep = records + .iter() + .all(|r| r.data.data.len() <= settings.max_record_size || settings.max_record_size == 0); + + if !keep { + counter!("atuin_record_too_large").increment(1); + + return Err( + ErrorResponse::reply("could not add records; record too large") + .with_status(StatusCode::BAD_REQUEST), + ); + } + + if let Err(e) = database.add_records(&user, &records).await { + error!("failed to add record: {}", e); + + return Err(ErrorResponse::reply("failed to add record") + .with_status(StatusCode::INTERNAL_SERVER_ERROR)); + } + + Ok(()) +} + +#[instrument(skip_all, fields(user.id = user.id.to_string()))] +pub(crate) async fn index( + UserAuth(user): UserAuth, + state: State<AppState>, +) -> Result<Json<RecordStatus>, ErrorResponseStatus<'static>> { + let State(AppState { + database, + settings: _, + }) = state; + + let record_index = match database.status(&user).await { + Ok(index) => index, + Err(e) => { + error!("failed to get record index: {}", e); + + return Err(ErrorResponse::reply("failed to calculate record index") + .with_status(StatusCode::INTERNAL_SERVER_ERROR)); + } + }; + + tracing::debug!(user = user.id.to_string(), "record index request"); + + Ok(Json(record_index)) +} + +#[derive(Deserialize)] +pub(crate) struct NextParams { + host: HostId, + tag: String, + start: Option<RecordIdx>, + count: u64, +} + +#[instrument(skip_all, fields(user.id = user.id.to_string()))] +pub(crate) async fn next( + params: Query<NextParams>, + UserAuth(user): UserAuth, + state: State<AppState>, +) -> Result<Json<Vec<Record<EncryptedData>>>, ErrorResponseStatus<'static>> { + let State(AppState { + database, + settings: _, + }) = state; + let params = params.0; + + let records = match database + .next_records(&user, params.host, params.tag, params.start, params.count) + .await + { + Ok(records) => records, + Err(e) => { + error!("failed to get record index: {}", e); + + return Err(ErrorResponse::reply("failed to calculate record index") + .with_status(StatusCode::INTERNAL_SERVER_ERROR)); + } + }; + + counter!("atuin_record_downloaded").increment(records.len() as u64); + + Ok(Json(records)) +} diff --git a/crates/server/src/lib.rs b/crates/server/src/lib.rs new file mode 100644 index 00000000..a4b10acf --- /dev/null +++ b/crates/server/src/lib.rs @@ -0,0 +1,86 @@ +use std::future::Future; +use std::net::SocketAddr; + +use axum::{Router, serve}; +use database::db::ServerPostgres; +use eyre::{Context, Result}; + +pub(crate) mod database; +mod handlers; +mod metrics; +mod router; + +pub(crate) use settings::Settings; + +pub(crate) mod settings; + +use tokio::net::TcpListener; +use tokio::signal; + +#[cfg(target_family = "unix")] +async fn shutdown_signal() { + let mut term = signal::unix::signal(signal::unix::SignalKind::terminate()) + .expect("failed to register signal handler"); + let mut interrupt = signal::unix::signal(signal::unix::SignalKind::interrupt()) + .expect("failed to register signal handler"); + + tokio::select! { + _ = term.recv() => {}, + _ = interrupt.recv() => {}, + }; + eprintln!("Shutting down gracefully..."); +} + +pub(crate) async fn launch(settings: Settings, addr: SocketAddr) -> Result<()> { + launch_with_tcp_listener( + settings, + TcpListener::bind(addr) + .await + .context("could not connect to socket")?, + shutdown_signal(), + ) + .await +} + +pub(crate) async fn launch_with_tcp_listener( + settings: Settings, + listener: TcpListener, + shutdown: impl Future<Output = ()> + Send + 'static, +) -> Result<()> { + let r = make_router(settings).await?; + + serve(listener, r.into_make_service()) + .with_graceful_shutdown(shutdown) + .await?; + + Ok(()) +} + +// The separate listener means it's much easier to ensure metrics are not accidentally exposed to +// the public. +pub(crate) async fn launch_metrics_server(host: String, port: u16) -> Result<()> { + let listener = TcpListener::bind((host, port)) + .await + .context("failed to bind metrics tcp")?; + + let recorder_handle = metrics::setup_metrics_recorder(); + + let router = Router::new().route( + "/metrics", + axum::routing::get(move || std::future::ready(recorder_handle.render())), + ); + + serve(listener, router.into_make_service()) + .with_graceful_shutdown(shutdown_signal()) + .await?; + + Ok(()) +} + +async fn make_router(settings: Settings) -> Result<Router, eyre::Error> { + let db = ServerPostgres::new(&settings.db_settings) + .await + .wrap_err_with(|| format!("failed to connect to db: {:?}", settings.db_settings))?; + let r = router::router(db, settings); + Ok(r) +} diff --git a/crates/server/src/metrics.rs b/crates/server/src/metrics.rs new file mode 100644 index 00000000..6380bef1 --- /dev/null +++ b/crates/server/src/metrics.rs @@ -0,0 +1,55 @@ +use std::time::Instant; + +use axum::{ + extract::{MatchedPath, Request}, + middleware::Next, + response::IntoResponse, +}; +use metrics_exporter_prometheus::{Matcher, PrometheusBuilder, PrometheusHandle}; + +pub(crate) fn setup_metrics_recorder() -> PrometheusHandle { + const EXPONENTIAL_SECONDS: &[f64] = &[ + 0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0, 10.0, + ]; + + PrometheusBuilder::new() + .set_buckets_for_metric( + Matcher::Full("http_requests_duration_seconds".to_string()), + EXPONENTIAL_SECONDS, + ) + .unwrap() + .install_recorder() + .unwrap() +} + +/// Middleware to record some common HTTP metrics +/// Generic over B to allow for arbitrary body types (eg Vec<u8>, Streams, a deserialized thing, etc) +/// Someday tower-http might provide a metrics middleware: <https://github.com/tower-rs/tower-http/issues/57> +pub(crate) async fn track_metrics(req: Request, next: Next) -> impl IntoResponse { + let start = Instant::now(); + + let path = req.extensions().get::<MatchedPath>().map_or_else( + || req.uri().path().to_owned(), + |matched_path| matched_path.as_str().to_owned(), + ); + + let method = req.method().clone(); + + // Run the rest of the request handling first, so we can measure it and get response + // codes. + let response = next.run(req).await; + + let latency = start.elapsed().as_secs_f64(); + let status = response.status().as_u16().to_string(); + + let labels = [ + ("method", method.to_string()), + ("path", path), + ("status", status), + ]; + + metrics::counter!("http_requests_total", &labels).increment(1); + metrics::histogram!("http_requests_duration_seconds", &labels).record(latency); + + response +} diff --git a/crates/server/src/router.rs b/crates/server/src/router.rs new file mode 100644 index 00000000..2a5c5f15 --- /dev/null +++ b/crates/server/src/router.rs @@ -0,0 +1,98 @@ +use crate::{ + atuin_common::api::{ATUIN_CARGO_VERSION, ATUIN_HEADER_VERSION, ErrorResponse}, + atuin_server::database::{db::ServerPostgres, models::User}, +}; +use axum::{ + Router, + extract::{FromRequestParts, Path, Request}, + http::{self, request::Parts}, + middleware::Next, + response::{IntoResponse, Response}, + routing::{get, post}, +}; +use eyre::Result; +use tower::ServiceBuilder; +use tower_http::trace::TraceLayer; +use uuid::Uuid; + +use super::handlers; +use crate::{ + handlers::{ErrorResponseStatus, RespExt}, + metrics, + settings::Settings, +}; + +pub(crate) struct UserAuth(pub(crate) User); + +impl FromRequestParts<AppState> for UserAuth { + type Rejection = ErrorResponseStatus<'static>; + + async fn from_request_parts( + req: &mut Parts, + state: &AppState, + ) -> Result<Self, Self::Rejection> { + let user_id = { + let Path(user_id) = + <Path<Uuid> as FromRequestParts<AppState>>::from_request_parts(req, state) + .await + .map_err(|_| { + ErrorResponse::reply("invalid user_id path param") + .with_status(http::StatusCode::BAD_REQUEST) + })?; + + user_id + }; + + let user = User { id: user_id }; + + Ok(Self(user)) + } +} + +async fn teapot() -> impl IntoResponse { + // This used to return 418: 🫖 + // Much as it was fun, it wasn't as useful or informative as it should be + (http::StatusCode::NOT_FOUND, "404 not found") +} + +/// Ensure that we only try and sync with clients on the same major version +async fn semver(request: Request, next: Next) -> Response { + let mut response = next.run(request).await; + response + .headers_mut() + .insert(ATUIN_HEADER_VERSION, ATUIN_CARGO_VERSION.parse().unwrap()); + + response +} + +#[derive(Clone)] +pub(crate) struct AppState { + pub(crate) database: ServerPostgres, + pub(crate) settings: Settings, +} + +pub(crate) fn router(database: ServerPostgres, settings: Settings) -> Router { + let routes = Router::new() + .route("/", get(handlers::index)) + .route("/api/v0/{user_id}/record", post(handlers::v0::record::post)) + .route("/api/v0/{user_id}/record", get(handlers::v0::record::index)) + .route( + "/api/v0/{user_id}/record/next", + get(handlers::v0::record::next), + ); + + let path = settings.path.as_str(); + if path.is_empty() { + routes + } else { + Router::new().nest(path, routes) + } + .fallback(teapot) + .with_state(AppState { database, settings }) + .layer( + ServiceBuilder::new() + .layer(TraceLayer::new_for_http()) + .layer(axum::middleware::from_fn(metrics::track_metrics)) + .layer(axum::middleware::from_fn(semver)), + ) +} diff --git a/crates/server/src/settings.rs b/crates/server/src/settings.rs new file mode 100644 index 00000000..6a32fb9b --- /dev/null +++ b/crates/server/src/settings.rs @@ -0,0 +1,96 @@ +use std::path::PathBuf; + +use config::{Config, Environment, File as ConfigFile, FileFormat}; +use eyre::{Result, eyre}; +use fs_err::create_dir_all; +use serde::{Deserialize, Serialize}; +use tracing::info; + +use crate::database::DbSettings; + +#[derive(Clone, Debug, Deserialize, Serialize)] +pub(crate) struct Metrics { + #[serde(alias = "enabled")] + pub(crate) enable: bool, + pub(crate) host: String, + pub(crate) port: u16, +} + +impl Default for Metrics { + fn default() -> Self { + Self { + enable: false, + host: String::from("127.0.0.1"), + port: 9001, + } + } +} + +#[derive(Clone, Debug, Deserialize, Serialize)] +pub(crate) struct Settings { + pub(crate) host: String, + pub(crate) port: u16, + pub(crate) path: String, + pub(crate) max_history_length: usize, + pub(crate) max_record_size: usize, + pub(crate) page_size: i64, + pub(crate) metrics: Metrics, + + /// Advertise a version that is not what we are _actually_ running + /// Many clients compare their version with api.atuin.sh, and if they differ, notify the user + /// that an update is available. + /// Now that we take beta releases, we should be able to advertise a different version to avoid + /// notifying users when the server runs something that is not a stable release. + pub(crate) fake_version: Option<String>, + + #[serde(flatten)] + #[expect(clippy::struct_field_names)] + pub(crate) db_settings: DbSettings, +} + +impl Settings { + pub(crate) fn new() -> Result<Self> { + // create the config file if it does not exist + let mut config_builder = Config::builder() + .set_default("host", "127.0.0.1")? + .set_default("port", 8888)? + .set_default("max_history_length", 8192)? + .set_default("max_record_size", 1024 * 1024 * 1024)? // pretty chonky + .set_default("path", "")? + .set_default("page_size", 1100)? + .set_default("metrics.enable", false)? + .set_default("metrics.host", "127.0.0.1")? + .set_default("metrics.port", 9001)? + .add_source( + Environment::with_prefix("atuin") + .prefix_separator("_") + .separator("__"), + ); + + if let Ok(mut config_file) = std::env::var("TURTLE_SERVER_CONFIG").map(PathBuf::from) { + config_builder = if config_file.exists() { + config_builder.add_source(ConfigFile::new( + config_file.to_str().unwrap(), + 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 + }; + } + + let config = config_builder.build()?; + + config + .try_deserialize() + .map_err(|e| eyre!("failed to deserialize: {}", e)) + } +} |