feat: change to sqlx

This commit is contained in:
kaichaosun 2026-06-11 16:25:13 +08:00
parent 1b20fc45a6
commit 9c572d65db
No known key found for this signature in database
GPG Key ID: 223E0F992F4F03BF
7 changed files with 1216 additions and 177 deletions

1081
Cargo.lock generated

File diff suppressed because it is too large Load Diff

View File

@ -18,10 +18,15 @@ base64 = "0.22"
clap = { version = "4", features = ["derive"] } clap = { version = "4", features = ["derive"] }
ed25519-dalek = "2.2.0" ed25519-dalek = "2.2.0"
hex = "0.4" hex = "0.4"
rusqlite = { version = "0.35", features = ["bundled-sqlcipher-vendored-openssl"] }
serde = { version = "1.0", features = ["derive"] } serde = { version = "1.0", features = ["derive"] }
serde_json = "1.0" serde_json = "1.0"
thiserror = "2" # SQLite via sqlx with a bundled libsqlite3 — no sqlcipher/OpenSSL, no system libs.
sqlx = { version = "0.8", default-features = false, features = [
"runtime-tokio",
"sqlite",
"macros",
"migrate",
] }
tokio = { version = "1", features = ["rt-multi-thread", "macros", "signal", "sync", "time"] } tokio = { version = "1", features = ["rt-multi-thread", "macros", "signal", "sync", "time"] }
tracing = "0.1" tracing = "0.1"
tracing-subscriber = { version = "0.3", features = ["env-filter"] } tracing-subscriber = { version = "0.3", features = ["env-filter"] }

View File

@ -5,13 +5,8 @@
######################################## ########################################
FROM rust:1-bookworm AS builder FROM rust:1-bookworm AS builder
# rusqlite's `bundled-sqlcipher-vendored-openssl` feature compiles SQLCipher and # sqlx's SQLite driver bundles libsqlite3 and compiles it with the C toolchain
# a vendored OpenSSL from source: the C toolchain ships in the base image, but # already in the base image — no sqlcipher, no OpenSSL, no extra apt packages.
# the OpenSSL build also needs perl + make.
RUN apt-get update \
&& apt-get install -y --no-install-recommends perl make \
&& rm -rf /var/lib/apt/lists/*
WORKDIR /app WORKDIR /app
# Build dependencies first against a stub binary so the (slow) SQLCipher/OpenSSL # Build dependencies first against a stub binary so the (slow) SQLCipher/OpenSSL

17
migrations/0001_init.sql Normal file
View File

@ -0,0 +1,17 @@
-- KeyPackage bundles: history per device, newest read back on fetch.
CREATE TABLE keypackages (
device_id TEXT NOT NULL,
received_at INTEGER NOT NULL,
payload BLOB NOT NULL,
signature BLOB NOT NULL,
PRIMARY KEY (device_id, received_at)
);
-- Account device-list bundles: exactly one row per account; newer upserts
-- replace the existing row (compare-and-swap on the payload's lamport).
CREATE TABLE account_bundles (
account_id TEXT NOT NULL PRIMARY KEY,
updated_at INTEGER NOT NULL,
payload BLOB NOT NULL,
signature BLOB NOT NULL
);

View File

@ -82,6 +82,7 @@ async fn submit(
signature: signature.to_vec(), signature: signature.to_vec(),
}, },
) )
.await
.map_err(ApiError::internal)?; .map_err(ApiError::internal)?;
Ok(StatusCode::NO_CONTENT) Ok(StatusCode::NO_CONTENT)
} }
@ -90,7 +91,11 @@ async fn fetch(
State(store): State<Arc<Store>>, State(store): State<Arc<Store>>,
Path(device_id): Path<String>, Path(device_id): Path<String>,
) -> Result<Json<FetchResponse>, ApiError> { ) -> Result<Json<FetchResponse>, ApiError> {
let Some(bundle) = store.latest(&device_id).map_err(ApiError::internal)? else { let Some(bundle) = store
.latest(&device_id)
.await
.map_err(ApiError::internal)?
else {
return Err(ApiError::not_found("no keypackage for device")); return Err(ApiError::not_found("no keypackage for device"));
}; };
Ok(Json(FetchResponse { Ok(Json(FetchResponse {
@ -172,6 +177,7 @@ async fn submit_account(
updated_at: 0, // filled in by store updated_at: 0, // filled in by store
}, },
) )
.await
.map_err(ApiError::internal)?; .map_err(ApiError::internal)?;
if !applied { if !applied {
return Err(ApiError::conflict( return Err(ApiError::conflict(
@ -189,7 +195,11 @@ async fn fetch_account(
State(store): State<Arc<Store>>, State(store): State<Arc<Store>>,
Path(account_id): Path<String>, Path(account_id): Path<String>,
) -> Result<Json<FetchAccountResponse>, ApiError> { ) -> Result<Json<FetchAccountResponse>, ApiError> {
let Some(bundle) = store.get_account(&account_id).map_err(ApiError::internal)? else { let Some(bundle) = store
.get_account(&account_id)
.await
.map_err(ApiError::internal)?
else {
return Err(ApiError::not_found("no account bundle for account_id")); return Err(ApiError::not_found("no account bundle for account_id"));
}; };
Ok(Json(FetchAccountResponse { Ok(Json(FetchAccountResponse {

View File

@ -57,7 +57,11 @@ async fn main() -> Result<()> {
let cli = Cli::parse(); let cli = Cli::parse();
let store = Arc::new(Store::open(&cli.db).context("failed to open store")?); let store = Arc::new(
Store::open(&cli.db)
.await
.context("failed to open store")?,
);
let prune_store = store.clone(); let prune_store = store.clone();
let max_per_id = cli.max_per_identity; let max_per_id = cli.max_per_identity;
@ -67,10 +71,10 @@ async fn main() -> Result<()> {
let mut ticker = tokio::time::interval(interval); let mut ticker = tokio::time::interval(interval);
loop { loop {
ticker.tick().await; ticker.tick().await;
if let Err(e) = prune_store.prune_key_packages(max_per_id, retention) { if let Err(e) = prune_store.prune_key_packages(max_per_id, retention).await {
tracing::warn!("prune (keypackages) failed: {e}"); tracing::warn!("prune (keypackages) failed: {e}");
} }
if let Err(e) = prune_store.prune_accounts(retention) { if let Err(e) = prune_store.prune_accounts(retention).await {
tracing::warn!("prune (accounts) failed: {e}"); tracing::warn!("prune (accounts) failed: {e}");
} }
} }

View File

@ -1,15 +1,15 @@
use std::path::Path; use std::path::Path;
use std::sync::Mutex;
use std::time::{Duration, SystemTime, UNIX_EPOCH}; use std::time::{Duration, SystemTime, UNIX_EPOCH};
use anyhow::{Context, Result}; use anyhow::{Context, Result};
use rusqlite::{Connection, OptionalExtension, params}; use sqlx::SqlitePool;
use sqlx::sqlite::{SqliteConnectOptions, SqliteJournalMode, SqlitePoolOptions};
pub struct Store { pub struct Store {
conn: Mutex<Connection>, pool: SqlitePool,
} }
#[derive(Debug, Clone)] #[derive(Debug, Clone, sqlx::FromRow)]
pub struct StoredKeyPackageBundle { pub struct StoredKeyPackageBundle {
/// The canonical signed payload, stored verbatim and returned as-is so /// The canonical signed payload, stored verbatim and returned as-is so
/// consumers verify over the exact bytes that were signed. /// consumers verify over the exact bytes that were signed.
@ -24,7 +24,7 @@ pub struct StoredKeyPackageBundle {
/// [`Store::upsert_account`]). `payload` is otherwise opaque to the server: it /// [`Store::upsert_account`]). `payload` is otherwise opaque to the server: it
/// encodes a lamport-timestamped list of device pubkeys signed by the account /// encodes a lamport-timestamped list of device pubkeys signed by the account
/// key so that consumers can verify the full device set. /// key so that consumers can verify the full device set.
#[derive(Debug, Clone)] #[derive(Debug, Clone, sqlx::FromRow)]
pub struct StoredAccountBundle { pub struct StoredAccountBundle {
/// The canonical signed payload, returned verbatim so consumers can verify /// The canonical signed payload, returned verbatim so consumers can verify
/// the account signature over the exact bytes. /// the account signature over the exact bytes.
@ -36,70 +36,73 @@ pub struct StoredAccountBundle {
} }
impl Store { impl Store {
pub fn open(path: &Path) -> Result<Self> { pub async fn open(path: &Path) -> Result<Self> {
// Create the db's parent directory if the caller pointed at a nested let is_memory = path == Path::new(":memory:");
// path (e.g. `tmp/registry.db`); SQLite won't create it and errors with
// "unable to open database file" otherwise. let mut options = SqliteConnectOptions::new()
if let Some(parent) = path.parent() .create_if_missing(true)
&& !parent.as_os_str().is_empty() .busy_timeout(Duration::from_secs(5));
{
std::fs::create_dir_all(parent) if is_memory {
.with_context(|| format!("create db directory {}", parent.display()))?; options = options.filename(":memory:");
} else {
// Create the db's parent directory if the caller pointed at a nested
// path (e.g. `tmp/registry.db`); SQLite won't create it and errors
// with "unable to open database file" otherwise.
if let Some(parent) = path.parent()
&& !parent.as_os_str().is_empty()
{
std::fs::create_dir_all(parent)
.with_context(|| format!("create db directory {}", parent.display()))?;
}
options = options.filename(path).journal_mode(SqliteJournalMode::Wal);
} }
let conn = Connection::open(path).context("open sqlite")?;
conn.execute_batch( // A shared in-memory database only lives as long as a connection is open,
"CREATE TABLE IF NOT EXISTS keypackages ( // so pin the `:memory:` pool to a single reused connection; file-backed
device_id TEXT NOT NULL, // pools can fan out across requests and the prune task.
received_at INTEGER NOT NULL, let pool = SqlitePoolOptions::new()
payload BLOB NOT NULL, .max_connections(if is_memory { 1 } else { 5 })
signature BLOB NOT NULL, .connect_with(options)
PRIMARY KEY (device_id, received_at) .await
); .context("open sqlite")?;
-- One row per account; newer upserts replace the existing row.
CREATE TABLE IF NOT EXISTS account_bundles ( sqlx::migrate!()
account_id TEXT NOT NULL PRIMARY KEY, .run(&pool)
updated_at INTEGER NOT NULL, .await
payload BLOB NOT NULL, .context("run migrations")?;
signature BLOB NOT NULL
);", Ok(Self { pool })
)?;
Ok(Self {
conn: Mutex::new(conn),
})
} }
pub fn insert(&self, device_id: &str, bundle: &StoredKeyPackageBundle) -> Result<()> { pub async fn insert(&self, device_id: &str, bundle: &StoredKeyPackageBundle) -> Result<()> {
let received_at = now_ms() as i64; let received_at = now_ms() as i64;
let conn = self.conn.lock().unwrap(); sqlx::query(
conn.execute( "INSERT INTO keypackages (device_id, received_at, payload, signature)
"INSERT INTO keypackages VALUES (?, ?, ?, ?)",
(device_id, received_at, payload, signature) )
VALUES (?1, ?2, ?3, ?4)", .bind(device_id)
params![device_id, received_at, bundle.payload, bundle.signature], .bind(received_at)
)?; .bind(&bundle.payload)
.bind(&bundle.signature)
.execute(&self.pool)
.await?;
Ok(()) Ok(())
} }
/// Returns the most recently received bundle for `device_id`. Scope A: the /// Returns the most recently received bundle for `device_id`. Scope A: the
/// chat layer consumes one bundle per device. When multi-keypackage fanout /// chat layer consumes one bundle per device. When multi-keypackage fanout
/// lands, switch this to return a `Vec<StoredKeyPackageBundle>`. /// lands, switch this to return a `Vec<StoredKeyPackageBundle>`.
pub fn latest(&self, device_id: &str) -> Result<Option<StoredKeyPackageBundle>> { pub async fn latest(&self, device_id: &str) -> Result<Option<StoredKeyPackageBundle>> {
let conn = self.conn.lock().unwrap(); let row = sqlx::query_as::<_, StoredKeyPackageBundle>(
let row = conn "SELECT payload, signature FROM keypackages
.query_row( WHERE device_id = ?
"SELECT payload, signature FROM keypackages ORDER BY received_at DESC
WHERE device_id = ?1 LIMIT 1",
ORDER BY received_at DESC )
LIMIT 1", .bind(device_id)
params![device_id], .fetch_optional(&self.pool)
|r| { .await?;
Ok(StoredKeyPackageBundle {
payload: r.get::<_, Vec<u8>>(0)?,
signature: r.get::<_, Vec<u8>>(1)?,
})
},
)
.optional()?;
Ok(row) Ok(row)
} }
@ -115,84 +118,85 @@ impl Store {
/// update so a replay can't keep a stale bundle alive past retention. /// update so a replay can't keep a stale bundle alive past retention.
/// ///
/// Returns `true` when the bundle was stored, `false` when it was rejected as /// Returns `true` when the bundle was stored, `false` when it was rejected as
/// stale. The compare-and-swap runs under the connection lock so concurrent /// stale. The compare-and-swap runs inside a transaction so a concurrent
/// publishes can't interleave a read with a write. The `updated_at` field of /// publish can't interleave the read with the write. The `updated_at` field of
/// `bundle` is ignored; the store stamps the row with the current time. /// `bundle` is ignored; the store stamps the row with the current time.
pub fn upsert_account( pub async fn upsert_account(
&self, &self,
account_id: &str, account_id: &str,
lamport: u64, lamport: u64,
bundle: &StoredAccountBundle, bundle: &StoredAccountBundle,
) -> Result<bool> { ) -> Result<bool> {
let updated_at = now_ms() as i64; let updated_at = now_ms() as i64;
let conn = self.conn.lock().unwrap();
let existing_lamport = conn let mut tx = self.pool.begin().await?;
.query_row( let existing_lamport = sqlx::query_scalar::<_, Vec<u8>>(
"SELECT payload FROM account_bundles WHERE account_id = ?1", "SELECT payload FROM account_bundles WHERE account_id = ?",
params![account_id], )
|r| r.get::<_, Vec<u8>>(0), .bind(account_id)
) .fetch_optional(&mut *tx)
.optional()? .await?
.and_then(|payload| payload_lamport(&payload)); .and_then(|payload| payload_lamport(&payload));
if let Some(stored) = existing_lamport if let Some(stored) = existing_lamport
&& lamport <= stored && lamport <= stored
{ {
// Dropping `tx` rolls the (read-only) transaction back.
return Ok(false); return Ok(false);
} }
conn.execute( sqlx::query(
"INSERT INTO account_bundles (account_id, updated_at, payload, signature) "INSERT INTO account_bundles (account_id, updated_at, payload, signature)
VALUES (?1, ?2, ?3, ?4) VALUES (?, ?, ?, ?)
ON CONFLICT(account_id) DO UPDATE SET ON CONFLICT(account_id) DO UPDATE SET
updated_at = excluded.updated_at, updated_at = excluded.updated_at,
payload = excluded.payload, payload = excluded.payload,
signature = excluded.signature", signature = excluded.signature",
params![account_id, updated_at, bundle.payload, bundle.signature], )
)?; .bind(account_id)
.bind(updated_at)
.bind(&bundle.payload)
.bind(&bundle.signature)
.execute(&mut *tx)
.await?;
tx.commit().await?;
Ok(true) Ok(true)
} }
/// Returns the stored bundle for `account_id`, or `None` if unknown. /// Returns the stored bundle for `account_id`, or `None` if unknown.
pub fn get_account(&self, account_id: &str) -> Result<Option<StoredAccountBundle>> { pub async fn get_account(&self, account_id: &str) -> Result<Option<StoredAccountBundle>> {
let conn = self.conn.lock().unwrap(); let row = sqlx::query_as::<_, StoredAccountBundle>(
let row = conn "SELECT payload, signature, updated_at FROM account_bundles
.query_row( WHERE account_id = ?",
"SELECT payload, signature, updated_at FROM account_bundles )
WHERE account_id = ?1", .bind(account_id)
params![account_id], .fetch_optional(&self.pool)
|r| { .await?;
Ok(StoredAccountBundle {
payload: r.get::<_, Vec<u8>>(0)?,
signature: r.get::<_, Vec<u8>>(1)?,
updated_at: r.get::<_, i64>(2)?,
})
},
)
.optional()?;
Ok(row) Ok(row)
} }
/// Drops account bundles that have not been refreshed within `retention`. /// Drops account bundles that have not been refreshed within `retention`.
pub fn prune_accounts(&self, retention: Duration) -> Result<()> { pub async fn prune_accounts(&self, retention: Duration) -> Result<()> {
let cutoff_ms = now_ms().saturating_sub(retention.as_millis() as u64) as i64; let cutoff_ms = now_ms().saturating_sub(retention.as_millis() as u64) as i64;
let conn = self.conn.lock().unwrap(); sqlx::query("DELETE FROM account_bundles WHERE updated_at < ?")
conn.execute( .bind(cutoff_ms)
"DELETE FROM account_bundles WHERE updated_at < ?1", .execute(&self.pool)
params![cutoff_ms], .await?;
)?;
Ok(()) Ok(())
} }
/// Drops bundles older than `retention` and keeps at most /// Drops bundles older than `retention` and keeps at most
/// `max_per_identity` per `device_id` — each device's history is bounded /// `max_per_identity` per `device_id` — each device's history is bounded
/// independently. /// independently.
pub fn prune_key_packages(&self, max_per_identity: usize, retention: Duration) -> Result<()> { pub async fn prune_key_packages(
&self,
max_per_identity: usize,
retention: Duration,
) -> Result<()> {
let cutoff_ms = now_ms().saturating_sub(retention.as_millis() as u64) as i64; let cutoff_ms = now_ms().saturating_sub(retention.as_millis() as u64) as i64;
let conn = self.conn.lock().unwrap(); sqlx::query("DELETE FROM keypackages WHERE received_at < ?")
conn.execute( .bind(cutoff_ms)
"DELETE FROM keypackages WHERE received_at < ?1", .execute(&self.pool)
params![cutoff_ms], .await?;
)?; sqlx::query(
conn.execute(
"DELETE FROM keypackages "DELETE FROM keypackages
WHERE rowid IN ( WHERE rowid IN (
SELECT rowid FROM ( SELECT rowid FROM (
@ -203,10 +207,12 @@ impl Store {
) AS rn ) AS rn
FROM keypackages FROM keypackages
) )
WHERE rn > ?1 WHERE rn > ?
)", )",
params![max_per_identity as i64], )
)?; .bind(max_per_identity as i64)
.execute(&self.pool)
.await?;
Ok(()) Ok(())
} }
} }
@ -259,40 +265,41 @@ mod tests {
} }
} }
fn upsert(store: &Store, account: &str, lamport: u64) -> bool { async fn upsert(store: &Store, account: &str, lamport: u64) -> bool {
store store
.upsert_account(account, lamport, &bundle(lamport)) .upsert_account(account, lamport, &bundle(lamport))
.await
.unwrap() .unwrap()
} }
#[test] #[tokio::test]
fn rejects_replayed_or_stale_lamport() { async fn rejects_replayed_or_stale_lamport() {
let store = Store::open(Path::new(":memory:")).unwrap(); let store = Store::open(Path::new(":memory:")).await.unwrap();
// First publish is always accepted. // First publish is always accepted.
assert!(upsert(&store, "acct", 5)); assert!(upsert(&store, "acct", 5).await);
// A strictly higher lamport replaces it. // A strictly higher lamport replaces it.
assert!(upsert(&store, "acct", 6)); assert!(upsert(&store, "acct", 6).await);
// Re-publishing the same lamport (a replay) is rejected. // Re-publishing the same lamport (a replay) is rejected.
assert!(!upsert(&store, "acct", 6)); assert!(!upsert(&store, "acct", 6).await);
// An older lamport (a downgrade) is rejected. // An older lamport (a downgrade) is rejected.
assert!(!upsert(&store, "acct", 4)); assert!(!upsert(&store, "acct", 4).await);
// The stored bundle is still the newest one accepted. // The stored bundle is still the newest one accepted.
let stored = store.get_account("acct").unwrap().unwrap(); let stored = store.get_account("acct").await.unwrap().unwrap();
assert_eq!(payload_lamport(&stored.payload), Some(6)); assert_eq!(payload_lamport(&stored.payload), Some(6));
} }
#[test] #[tokio::test]
fn stale_publish_does_not_refresh_retention_clock() { async fn stale_publish_does_not_refresh_retention_clock() {
let store = Store::open(Path::new(":memory:")).unwrap(); let store = Store::open(Path::new(":memory:")).await.unwrap();
assert!(upsert(&store, "acct", 9)); assert!(upsert(&store, "acct", 9).await);
let after_first = store.get_account("acct").unwrap().unwrap().updated_at; let after_first = store.get_account("acct").await.unwrap().unwrap().updated_at;
// A rejected (stale) publish must not bump updated_at, so a replay can't // A rejected (stale) publish must not bump updated_at, so a replay can't
// keep a stale bundle alive past the retention window. // keep a stale bundle alive past the retention window.
assert!(!upsert(&store, "acct", 9)); assert!(!upsert(&store, "acct", 9).await);
let after_replay = store.get_account("acct").unwrap().unwrap().updated_at; let after_replay = store.get_account("acct").await.unwrap().unwrap().updated_at;
assert_eq!(after_first, after_replay); assert_eq!(after_first, after_replay);
} }