diff options
Diffstat (limited to 'crates/server/src')
| -rw-r--r-- | crates/server/src/database/db/mod.rs | 274 | ||||
| -rw-r--r-- | crates/server/src/database/db/wrappers.rs | 32 | ||||
| -rw-r--r-- | crates/server/src/database/mod.rs | 99 | ||||
| -rw-r--r-- | crates/server/src/database/models.rs | 5 | ||||
| -rw-r--r-- | crates/server/src/handlers/mod.rs | 54 | ||||
| -rw-r--r-- | crates/server/src/handlers/v0/mod.rs | 1 | ||||
| -rw-r--r-- | crates/server/src/handlers/v0/record.rs | 113 | ||||
| -rw-r--r-- | crates/server/src/lib.rs | 86 | ||||
| -rw-r--r-- | crates/server/src/metrics.rs | 55 | ||||
| -rw-r--r-- | crates/server/src/router.rs | 98 | ||||
| -rw-r--r-- | crates/server/src/settings.rs | 96 |
11 files changed, 913 insertions, 0 deletions
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)) + } +} |
