Menu
AkurAI-Build
publicLatest change e5121cbf4cf5b44c00928705f6f3f3988d273833 - server: surface vector capacity/validation errors in the 400 body (closes #7) by Olafur Bui
use std::{
env,
error::Error,
fmt,
fs::{self, File, OpenOptions},
io::{Read, Seek, SeekFrom, Write},
path::{Path, PathBuf},
sync::{Arc, Mutex, MutexGuard},
time::{Duration, SystemTime, UNIX_EPOCH},
};
use anyhow::{Context, Result, anyhow, ensure};
use rusqlite::{
Connection, OpenFlags, TransactionBehavior, config::DbConfig, limits::Limit, params,
};
use serde::Serialize;
use subtle::ConstantTimeEq;
use zeroize::Zeroizing;
const APPLICATION_ID: i64 = 0x424E_464B;
const MAX_CONTENT_BYTES: usize = 64 * 1024;
const MAX_DIMENSIONS: usize = 4096;
const MAX_CANDIDATES: usize = 10_000;
const MAX_SEARCH_WORK: usize = 10_000_000;
const MAX_RESULTS: usize = 25;
const MIGRATION_LEDGER_SQL: &str = "CREATE TABLE IF NOT EXISTS _migrations (
version INTEGER PRIMARY KEY,
name TEXT NOT NULL UNIQUE,
sql TEXT NOT NULL,
applied_at INTEGER NOT NULL DEFAULT (unixepoch())
) STRICT;";
const VECTORS_TABLE_SQL: &str = "CREATE TABLE vectors (
tenant_id TEXT NOT NULL CHECK (length(tenant_id) BETWEEN 1 AND 128),
model TEXT NOT NULL CHECK (length(model) BETWEEN 1 AND 128),
id TEXT NOT NULL CHECK (length(id) BETWEEN 1 AND 128),
content TEXT NOT NULL CHECK (length(content) <= 65536),
embedding BLOB NOT NULL CHECK (length(embedding) = dimensions * 4),
dimensions INTEGER NOT NULL CHECK (dimensions BETWEEN 1 AND 4096),
updated_at INTEGER NOT NULL DEFAULT (unixepoch()),
PRIMARY KEY (tenant_id, model, id)
) STRICT;";
const VECTORS_INDEX_SQL: &str =
"CREATE INDEX vectors_search ON vectors(tenant_id, model, dimensions);";
const MIGRATIONS: &[(i64, &str, &str)] = &[(
1,
"initial vectors schema",
include_str!("../migrations/001_init.sql"),
)];
#[derive(Clone)]
pub struct Database {
connection: Arc<Mutex<Connection>>,
_lock: Option<Arc<File>>,
}
#[derive(Debug, Serialize)]
pub struct VectorMatch {
pub id: String,
pub content: String,
pub score: f32,
}
#[derive(Debug)]
pub struct VectorInputError(pub(crate) String);
impl fmt::Display for VectorInputError {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter.write_str(&self.0)
}
}
impl Error for VectorInputError {}
impl Database {
pub fn open_existing(path: &Path, key: &str) -> Result<Self> {
validate_secret(key)?;
ensure!(
metadata_if_exists(path)?.is_some(),
"no database at {}; run `bunfork migrate` to create it",
path.display()
);
validate_existing_file(path, "database")?;
ensure_secure_parent(path)?;
let lock = Arc::new(open_lock(path)?);
lock.lock_shared()
.context("database is locked for maintenance")?;
let database = Self {
connection: Arc::new(Mutex::new(open_connection(path, key, true, false)?)),
_lock: Some(lock),
};
database.validate_schema()?;
Ok(database)
}
pub fn open_for_migration(path: &Path, key: &str) -> Result<Self> {
validate_secret(key)?;
prepare_database_path(path)?;
let lock = Arc::new(open_lock(path)?);
lock.try_lock()
.context("database is in use; stop the server before migrating")?;
let connection = open_connection(path, key, true, true)?;
Ok(Self {
connection: Arc::new(Mutex::new(connection)),
_lock: Some(lock),
})
}
#[cfg(test)]
pub(crate) fn memory(key: &str) -> Result<Self> {
validate_secret(key)?;
let connection =
Connection::open_in_memory().context("open in-memory SQLCipher database")?;
configure_connection(&connection, key, false)?;
Ok(Self {
connection: Arc::new(Mutex::new(connection)),
_lock: None,
})
}
pub fn migrate(&self) -> Result<usize> {
let mut connection = self.connection()?;
let transaction = connection
.transaction_with_behavior(TransactionBehavior::Immediate)
.context("start migration transaction")?;
transaction.execute_batch(MIGRATION_LEDGER_SQL)?;
let existing = migration_history(&transaction)?;
validate_migration_prefix(&existing)?;
let mut applied = 0;
for (version, name, sql) in &MIGRATIONS[existing.len()..] {
transaction
.execute_batch(sql)
.with_context(|| format!("apply migration {version}: {name}"))?;
transaction.execute(
"INSERT INTO _migrations(version, name, sql) VALUES (?1, ?2, ?3)",
params![version, name, sql],
)?;
// Stamp the version pragmas programmatically so a future
// migration that omits its own PRAGMA lines cannot leave the
// database failing schema validation after a clean migrate.
transaction.pragma_update(None, "application_id", APPLICATION_ID)?;
transaction.pragma_update(None, "user_version", *version)?;
applied += 1;
}
transaction.commit().context("commit migrations")?;
Ok(applied)
}
pub fn validate_schema(&self) -> Result<()> {
let connection = self.connection()?;
validate_application_schema(&connection)
}
pub fn check_ready(&self) -> Result<()> {
self.connection()?.query_row("SELECT 1", [], |_| Ok(()))?;
Ok(())
}
pub fn upsert_vector(
&self,
tenant: &str,
model: &str,
id: &str,
content: &str,
embedding: &[f32],
) -> Result<()> {
validate_vector_upsert(tenant, model, id, content, embedding)?;
let normalized = normalize_embedding(embedding)?;
let dimensions = i64::try_from(normalized.len()).context("embedding is too large")?;
let blob = encode_embedding(&normalized);
let mut connection = self.connection()?;
let transaction = connection.transaction_with_behavior(TransactionBehavior::Immediate)?;
// Enforce the search work budget at write time: without this cap a
// tenant could insert past the budget and permanently break search
// for that tenant/model/dimension with no recovery besides deletes.
let capacity = vector_capacity(normalized.len());
let existing: i64 = transaction.query_row(
"SELECT count(*) FROM vectors
WHERE tenant_id = ?1 AND model = ?2 AND dimensions = ?3 AND id != ?4",
params![tenant, model, dimensions, id],
|row| row.get(0),
)?;
if usize::try_from(existing).context("invalid vector count")? >= capacity {
return Err(VectorInputError(format!(
"vector capacity ({capacity}) reached for this tenant, model, and \
dimension; delete vectors before adding more"
))
.into());
}
transaction.execute(
"INSERT INTO vectors(tenant_id, model, id, content, embedding, dimensions)
VALUES (?1, ?2, ?3, ?4, ?5, ?6)
ON CONFLICT(tenant_id, model, id) DO UPDATE SET
content = excluded.content,
embedding = excluded.embedding,
dimensions = excluded.dimensions,
updated_at = unixepoch()",
params![tenant, model, id, content, blob, dimensions],
)?;
transaction.commit()?;
Ok(())
}
pub fn search_vectors(
&self,
tenant: &str,
model: &str,
embedding: &[f32],
limit: usize,
) -> Result<Vec<VectorMatch>> {
validate_vector_search(tenant, model, embedding, limit)?;
let query = normalize_embedding(embedding)?;
let dimensions = i64::try_from(query.len()).context("embedding is too large")?;
let mut connection = self.connection()?;
let transaction = connection.transaction_with_behavior(TransactionBehavior::Deferred)?;
let work_limit = vector_capacity(query.len());
let candidate_count: i64 = transaction.query_row(
"SELECT count(*) FROM vectors
WHERE tenant_id = ?1 AND model = ?2 AND dimensions = ?3",
params![tenant, model, dimensions],
|row| row.get(0),
)?;
let candidate_count = usize::try_from(candidate_count).context("invalid vector count")?;
if candidate_count > work_limit {
return Err(VectorInputError(
"vector search budget exceeded for this tenant, model, and dimension".to_owned(),
)
.into());
}
let mut statement = transaction.prepare(
"SELECT id, embedding
FROM vectors
WHERE tenant_id = ?1 AND model = ?2 AND dimensions = ?3
LIMIT ?4",
)?;
let candidate_limit = i64::try_from(work_limit + 1)?;
let rows = statement
.query_map(params![tenant, model, dimensions, candidate_limit], |row| {
Ok((row.get::<_, String>(0)?, row.get::<_, Vec<u8>>(1)?))
})?;
// ponytail: exact O(N*d), bounded to ten million multiply-adds and O(k) memory.
// Add a statically linked ANN backend only after a measured corpus needs it.
let mut winners: Vec<(String, f32)> = Vec::with_capacity(limit);
for row in rows {
let (id, blob) = row?;
ensure!(
blob.len() == query.len() * 4,
"stored vector dimensions are corrupt"
);
// Compute the dot product directly over the stored blob: this
// avoids one Vec<f32> allocation and a redundant finiteness pass
// per candidate (writes already validate finiteness).
let score = blob
.chunks_exact(4)
.zip(&query)
.map(|(bytes, right)| {
f32::from_le_bytes([bytes[0], bytes[1], bytes[2], bytes[3]]) * right
})
.sum();
if winners.len() < limit {
winners.push((id, score));
continue;
}
let mut worst = 0;
for index in 1..winners.len() {
if is_better(&winners[worst], &winners[index]) {
worst = index;
}
}
let candidate = (id, score);
if is_better(&candidate, &winners[worst]) {
winners[worst] = candidate;
}
}
winners.sort_by(|left, right| {
right
.1
.total_cmp(&left.1)
.then_with(|| left.0.cmp(&right.0))
});
drop(statement);
let mut content = transaction.prepare(
"SELECT content FROM vectors
WHERE tenant_id = ?1 AND model = ?2 AND id = ?3",
)?;
let matches = winners
.into_iter()
.map(|(id, score)| {
let value = content.query_row(params![tenant, model, id], |row| row.get(0))?;
Ok(VectorMatch {
id,
content: value,
score,
})
})
.collect::<rusqlite::Result<Vec<_>>>()?;
drop(content);
transaction.commit()?;
Ok(matches)
}
pub fn delete_vector(&self, tenant: &str, model: &str, id: &str) -> Result<bool> {
validate_vector_id(tenant, model, id)?;
let changed = self.connection()?.execute(
"DELETE FROM vectors WHERE tenant_id = ?1 AND model = ?2 AND id = ?3",
params![tenant, model, id],
)?;
Ok(changed == 1)
}
pub fn backup(&self, destination: &Path, key: &str) -> Result<()> {
validate_secret(key)?;
self.validate_schema()?;
ensure!(
!destination.exists(),
"backup already exists: {}",
destination.display()
);
prepare_output_parent(destination)?;
let temporary = sibling_path(destination, &format!(".tmp-{}", std::process::id()))?;
ensure!(!temporary.exists(), "temporary backup path already exists");
let temporary_file = open_new_private(&temporary)?;
crate::fsguard::ensure_path_matches_file(&temporary, &temporary_file, "temporary backup")?;
let mut published = false;
let result = self.export_locked(&temporary, key).and_then(|()| {
crate::fsguard::ensure_path_matches_file(
&temporary,
&temporary_file,
"temporary backup",
)?;
let verification = open_read_only(&temporary, key)?;
validate_application_schema(&verification)?;
drop(verification);
crate::fsguard::ensure_path_matches_file(
&temporary,
&temporary_file,
"temporary backup",
)?;
temporary_file.sync_all()?;
fs::hard_link(&temporary, destination)
.with_context(|| format!("publish backup {}", destination.display()))?;
published = true;
crate::fsguard::ensure_path_matches_file(
destination,
&temporary_file,
"published backup",
)?;
fs::remove_file(&temporary)?;
let verification = open_read_only(destination, key)?;
validate_application_schema(&verification)?;
drop(verification);
crate::fsguard::ensure_path_matches_file(
destination,
&temporary_file,
"published backup",
)?;
sync_parent(destination)
});
if result.is_err() {
remove_file_if_matches(&temporary, &temporary_file);
if published {
remove_file_if_matches(destination, &temporary_file);
}
}
result
}
fn export_locked(&self, destination: &Path, key: &str) -> Result<()> {
let connection = self.connection()?;
validate_application_schema(&connection)?;
let destination = destination
.to_str()
.context("backup path must contain valid UTF-8")?;
let key = decode_secret(key)?;
connection.set_db_config(DbConfig::SQLITE_DBCONFIG_ENABLE_ATTACH_CREATE, true)?;
let export = (|| -> Result<()> {
connection.set_db_config(DbConfig::SQLITE_DBCONFIG_ENABLE_ATTACH_WRITE, true)?;
connection.execute_batch("BEGIN IMMEDIATE")?;
connection.execute(
"ATTACH DATABASE ?1 AS bunfork_backup KEY ?2",
params![destination, key.as_slice()],
)?;
connection.execute_batch("SELECT sqlcipher_export('bunfork_backup')")?;
connection.execute_batch(&format!(
"PRAGMA bunfork_backup.application_id = {APPLICATION_ID};
PRAGMA bunfork_backup.user_version = {};",
MIGRATIONS.len()
))?;
connection.execute_batch("COMMIT")?;
connection.execute_batch("DETACH DATABASE bunfork_backup")?;
Ok(())
})();
if export.is_err() {
let _ = connection.execute_batch("ROLLBACK");
let _ = connection.execute_batch("DETACH DATABASE bunfork_backup");
}
let disable_write =
connection.set_db_config(DbConfig::SQLITE_DBCONFIG_ENABLE_ATTACH_WRITE, false);
let disable_create =
connection.set_db_config(DbConfig::SQLITE_DBCONFIG_ENABLE_ATTACH_CREATE, false);
export?;
disable_write?;
disable_create?;
Ok(())
}
fn connection(&self) -> Result<MutexGuard<'_, Connection>> {
self.connection
.lock()
.map_err(|_| anyhow!("database lock is poisoned"))
}
}
pub fn validate_vector_upsert(
tenant: &str,
model: &str,
id: &str,
content: &str,
embedding: &[f32],
) -> Result<()> {
validate_vector_id(tenant, model, id).map_err(vector_input)?;
if content.len() > MAX_CONTENT_BYTES {
return Err(VectorInputError(format!("content exceeds {MAX_CONTENT_BYTES} bytes")).into());
}
normalize_embedding(embedding).map_err(vector_input)?;
Ok(())
}
pub fn validate_vector_search(
tenant: &str,
model: &str,
embedding: &[f32],
limit: usize,
) -> Result<()> {
validate_identifier("tenant", tenant).map_err(vector_input)?;
validate_identifier("model", model).map_err(vector_input)?;
if !(1..=MAX_RESULTS).contains(&limit) {
return Err(VectorInputError(format!("limit must be between 1 and {MAX_RESULTS}")).into());
}
normalize_embedding(embedding).map_err(vector_input)?;
Ok(())
}
pub fn validate_vector_id(tenant: &str, model: &str, id: &str) -> Result<()> {
validate_identifier("tenant", tenant).map_err(vector_input)?;
validate_identifier("model", model).map_err(vector_input)?;
validate_identifier("vector id", id).map_err(vector_input)
}
fn vector_input(error: anyhow::Error) -> anyhow::Error {
VectorInputError(error.to_string()).into()
}
pub fn load_secret(path: &Path, environment: &str) -> Result<Zeroizing<String>> {
let value = if metadata_if_exists(path)?.is_some() {
validate_existing_file(path, "secret file")?;
ensure_secure_parent(path)?;
let file = crate::fsguard::open_nofollow(path, false)?;
let metadata = file.metadata()?;
ensure!(metadata.len() <= 1024, "secret file is unexpectedly large");
ensure_private_mode(path, &metadata)?;
let mut file = file;
// Read into a zeroizing buffer with enough capacity to avoid
// reallocation, then trim in place: no intermediate copy of the
// secret is ever dropped unwiped.
let mut value = Zeroizing::new(String::with_capacity(4096));
file.read_to_string(&mut value)?;
while value.ends_with(|character: char| character.is_whitespace()) {
value.pop();
}
let leading = value.len() - value.trim_start().len();
value.drain(..leading);
value
} else {
Zeroizing::new(env::var(environment).with_context(|| {
format!(
"missing {}; create {} with `bunfork keygen --out {}`",
environment,
path.display(),
path.display()
)
})?)
};
validate_secret(&value)?;
Ok(value)
}
pub fn generate_secret(path: &Path) -> Result<()> {
ensure!(
metadata_if_exists(path)?.is_none(),
"refusing to overwrite {}",
path.display()
);
if let Some(parent) = path
.parent()
.filter(|parent| !parent.as_os_str().is_empty())
&& !parent.exists()
{
create_private_dir(parent)?;
}
ensure_secure_parent(path)?;
let mut bytes = Zeroizing::new([0_u8; 32]);
getrandom::fill(bytes.as_mut()).context("read operating-system randomness")?;
let mut encoded = Zeroizing::new(String::with_capacity(65));
for byte in bytes.iter() {
use std::fmt::Write as _;
write!(encoded, "{byte:02x}").context("encode secret")?;
}
encoded.push('\n');
let mut file = open_new_private(path)?;
file.write_all(encoded.as_bytes())?;
file.sync_all()?;
Ok(())
}
pub fn secrets_equal(left: &str, right: &str) -> Result<bool> {
let left = decode_secret(left)?;
let right = decode_secret(right)?;
Ok(bool::from(left.as_slice().ct_eq(right.as_slice())))
}
pub fn parse_embedding(value: &str) -> Result<Vec<f32>> {
let embedding = value
.split(',')
.map(str::trim)
.map(|part| {
ensure!(!part.is_empty(), "embedding contains an empty value");
part.parse::<f32>()
.with_context(|| format!("invalid embedding value: {part}"))
})
.collect::<Result<Vec<_>>>()?;
normalize_embedding(&embedding).map(|_| embedding)
}
pub fn restore_database(live: &Path, backup: &Path, key: &str) -> Result<PathBuf> {
validate_secret(key)?;
ensure_secure_parent(live)?;
ensure_secure_parent(backup)?;
let mut backup_file = open_validated_existing_file(backup, "backup", false)?;
let backup_connection = open_read_only(backup, key)?;
validate_application_schema(&backup_connection)?;
drop(backup_connection);
crate::fsguard::ensure_path_matches_file(backup, &backup_file, "backup")?;
let lock = open_lock(live)?;
lock.try_lock()
.context("database is in use; stop the server before restoring")?;
let mut live_file = open_validated_existing_file(live, "live database", false)?;
ensure_files_differ(
live,
&live_file,
backup,
&backup_file,
"backup and live database must differ",
)?;
let live_connection = open_connection(live, key, true, false)?;
validate_application_schema(&live_connection)?;
let (busy, _log_pages, _checkpointed_pages): (i64, i64, i64) =
live_connection.query_row("PRAGMA wal_checkpoint(TRUNCATE)", [], |row| {
Ok((row.get(0)?, row.get(1)?, row.get(2)?))
})?;
ensure!(busy == 0, "live database WAL is busy; restore aborted");
let journal_mode: String =
live_connection
.pragma_update_and_check(None, "journal_mode", "DELETE", |row| row.get(0))?;
ensure!(
journal_mode.eq_ignore_ascii_case("delete"),
"failed to leave WAL mode before restore"
);
drop(live_connection);
crate::fsguard::ensure_path_matches_file(live, &live_file, "live database")?;
ensure_no_sidecars(live)?;
let temporary = sibling_path(live, &format!(".restore-{}", std::process::id()))?;
ensure!(
fs::symlink_metadata(&temporary).is_err(),
"temporary restore path already exists"
);
let mut temporary_file = open_new_private(&temporary)?;
crate::fsguard::ensure_path_matches_file(
&temporary,
&temporary_file,
"temporary restore database",
)?;
let result = (|| -> Result<PathBuf> {
copy_file_contents(&mut backup_file, &mut temporary_file)
.context("copy verified backup for restore")?;
crate::fsguard::ensure_path_matches_file(
&temporary,
&temporary_file,
"temporary restore database",
)?;
let staged = open_read_only(&temporary, key)?;
validate_application_schema(&staged)?;
drop(staged);
crate::fsguard::ensure_path_matches_file(
&temporary,
&temporary_file,
"temporary restore database",
)?;
let old = sibling_path(live, &format!(".pre-restore-{}", unix_timestamp()?))?;
ensure!(
fs::symlink_metadata(&old).is_err(),
"pre-restore database path already exists"
);
let mut old_file = open_new_private(&old)?;
copy_file_contents(&mut live_file, &mut old_file)
.context("preserve pre-restore database")?;
crate::fsguard::ensure_path_matches_file(&old, &old_file, "pre-restore database")?;
let preserved = open_read_only(&old, key)?;
validate_application_schema(&preserved)?;
drop(preserved);
crate::fsguard::ensure_path_matches_file(&old, &old_file, "pre-restore database")?;
sync_parent(&old)?;
crate::fsguard::ensure_path_matches_file(live, &live_file, "live database")?;
crate::fsguard::ensure_path_matches_file(
&temporary,
&temporary_file,
"temporary restore database",
)?;
fs::rename(&temporary, live).context("atomically install restored database")?;
crate::fsguard::ensure_path_matches_file(live, &temporary_file, "restored database")?;
let installed = open_read_only(live, key)?;
validate_application_schema(&installed)?;
drop(installed);
crate::fsguard::ensure_path_matches_file(live, &temporary_file, "restored database")?;
sync_parent(live)?;
Ok(old)
})();
if result.is_err() {
remove_file_if_matches(&temporary, &temporary_file);
}
result
}
fn open_connection(path: &Path, key: &str, wal: bool, create: bool) -> Result<Connection> {
let mut flags = OpenFlags::SQLITE_OPEN_READ_WRITE
| OpenFlags::SQLITE_OPEN_FULL_MUTEX
| OpenFlags::SQLITE_OPEN_NOFOLLOW;
if create {
flags |= OpenFlags::SQLITE_OPEN_CREATE;
}
let connection = Connection::open_with_flags(path, flags)
.with_context(|| format!("open encrypted database {}", path.display()))?;
configure_connection(&connection, key, wal)?;
Ok(connection)
}
fn open_read_only(path: &Path, key: &str) -> Result<Connection> {
validate_existing_file(path, "backup")?;
ensure_secure_parent(path)?;
let connection = Connection::open_with_flags(
path,
OpenFlags::SQLITE_OPEN_READ_ONLY
| OpenFlags::SQLITE_OPEN_FULL_MUTEX
| OpenFlags::SQLITE_OPEN_NOFOLLOW,
)
.with_context(|| format!("open encrypted database {} read-only", path.display()))?;
configure_connection(&connection, key, false)?;
connection.pragma_update(None, "query_only", "ON")?;
Ok(connection)
}
fn configure_connection(connection: &Connection, key: &str, wal: bool) -> Result<()> {
apply_key(connection, key)?;
let cipher_version: String =
connection.pragma_query_value(None, "cipher_version", |row| row.get(0))?;
ensure!(
!cipher_version.is_empty(),
"SQLCipher support is unavailable"
);
let cipher_status: String =
connection.pragma_query_value(None, "cipher_status", |row| row.get(0))?;
ensure!(cipher_status == "1", "SQLCipher key is not active");
connection
.query_row("SELECT count(*) FROM sqlite_schema", [], |_| Ok(()))
.context("database key is incorrect or database is corrupt")?;
connection.busy_timeout(Duration::from_secs(5))?;
connection.pragma_update(None, "cipher_memory_security", "ON")?;
connection.pragma_update(None, "foreign_keys", "ON")?;
connection.pragma_update(None, "synchronous", "FULL")?;
connection.pragma_update(None, "temp_store", "MEMORY")?;
connection.pragma_update(None, "secure_delete", "ON")?;
connection.pragma_update(None, "mmap_size", 0_i64)?;
connection.pragma_update(None, "cell_size_check", "ON")?;
if wal {
let mode: String =
connection.pragma_update_and_check(None, "journal_mode", "WAL", |row| row.get(0))?;
ensure!(
mode.eq_ignore_ascii_case("wal"),
"failed to enable WAL mode"
);
}
connection.set_db_config(DbConfig::SQLITE_DBCONFIG_DEFENSIVE, true)?;
connection.set_db_config(DbConfig::SQLITE_DBCONFIG_WRITABLE_SCHEMA, false)?;
connection.set_db_config(DbConfig::SQLITE_DBCONFIG_DQS_DDL, false)?;
connection.set_db_config(DbConfig::SQLITE_DBCONFIG_DQS_DML, false)?;
connection.set_db_config(DbConfig::SQLITE_DBCONFIG_TRUSTED_SCHEMA, false)?;
connection.set_db_config(DbConfig::SQLITE_DBCONFIG_ENABLE_FTS3_TOKENIZER, false)?;
connection.set_db_config(DbConfig::SQLITE_DBCONFIG_ENABLE_ATTACH_CREATE, false)?;
connection.set_db_config(DbConfig::SQLITE_DBCONFIG_ENABLE_ATTACH_WRITE, false)?;
connection.set_limit(Limit::SQLITE_LIMIT_LENGTH, 8 * 1024 * 1024)?;
connection.set_limit(Limit::SQLITE_LIMIT_SQL_LENGTH, 256 * 1024)?;
connection.set_limit(Limit::SQLITE_LIMIT_COLUMN, 64)?;
connection.set_limit(Limit::SQLITE_LIMIT_EXPR_DEPTH, 50)?;
connection.set_limit(Limit::SQLITE_LIMIT_COMPOUND_SELECT, 10)?;
connection.set_limit(Limit::SQLITE_LIMIT_FUNCTION_ARG, 32)?;
connection.set_limit(Limit::SQLITE_LIMIT_ATTACHED, 1)?;
connection.set_limit(Limit::SQLITE_LIMIT_LIKE_PATTERN_LENGTH, 4096)?;
connection.set_limit(Limit::SQLITE_LIMIT_VARIABLE_NUMBER, 100)?;
connection.set_limit(Limit::SQLITE_LIMIT_TRIGGER_DEPTH, 8)?;
connection.set_limit(Limit::SQLITE_LIMIT_WORKER_THREADS, 2)?;
Ok(())
}
#[allow(unsafe_code)]
fn apply_key(connection: &Connection, encoded: &str) -> Result<()> {
let key = decode_secret(encoded)?;
let length = i32::try_from(key.len()).context("database key is too long")?;
// SAFETY: rusqlite owns a live sqlite3 handle for this call, and SQLCipher copies the
// pointed-to bytes synchronously before the zeroizing buffer is dropped.
let result = unsafe {
rusqlite::ffi::sqlite3_key(
connection.handle(),
key.as_ptr().cast::<std::ffi::c_void>(),
length,
)
};
ensure!(
result == rusqlite::ffi::SQLITE_OK,
"SQLCipher rejected the database key"
);
Ok(())
}
fn verify_connection(connection: &Connection) -> Result<()> {
let quick_check: String =
connection.pragma_query_value(None, "quick_check", |row| row.get(0))?;
ensure!(
quick_check == "ok",
"SQLite quick check failed: {quick_check}"
);
let database_file: String = connection.query_row(
"SELECT file FROM pragma_database_list WHERE name = 'main'",
[],
|row| row.get(0),
)?;
if !database_file.is_empty() {
let mut statement = connection.prepare("PRAGMA cipher_integrity_check")?;
let mut rows = statement.query([])?;
if let Some(row) = rows.next()? {
let result: String = row.get(0)?;
ensure!(
result.is_empty() || result == "ok",
"SQLCipher integrity check failed: {result}"
);
}
}
Ok(())
}
type MigrationRecord = (i64, String, String);
fn migration_history(connection: &Connection) -> Result<Vec<MigrationRecord>> {
let mut statement = connection
.prepare("SELECT version, name, sql FROM _migrations ORDER BY version")
.context("database is missing the Bunfork migration ledger")?;
statement
.query_map([], |row| {
Ok((
row.get::<_, i64>(0)?,
row.get::<_, String>(1)?,
row.get::<_, String>(2)?,
))
})?
.collect::<rusqlite::Result<Vec<_>>>()
.map_err(Into::into)
}
fn validate_migration_prefix(existing: &[MigrationRecord]) -> Result<()> {
ensure!(
existing.len() <= MIGRATIONS.len(),
"database contains migrations newer than this bunfork binary"
);
for (index, (version, name, sql)) in existing.iter().enumerate() {
let expected = MIGRATIONS[index];
ensure!(
(*version, name.as_str(), sql.as_str()) == expected,
"migration history drift at version {version}"
);
}
Ok(())
}
fn validate_application_schema(connection: &Connection) -> Result<()> {
verify_connection(connection)?;
let application_id: i64 =
connection.pragma_query_value(None, "application_id", |row| row.get(0))?;
ensure!(
application_id == APPLICATION_ID,
"database is not a Bunfork database"
);
let user_version: i64 =
connection.pragma_query_value(None, "user_version", |row| row.get(0))?;
ensure!(
user_version == i64::try_from(MIGRATIONS.len())?,
"database schema version does not match this Bunfork binary"
);
let history = migration_history(connection)?;
validate_migration_prefix(&history)?;
ensure!(
history.len() == MIGRATIONS.len(),
"database has unapplied migrations; run `bunfork migrate`"
);
let table_count: i64 = connection.query_row(
"SELECT count(*) FROM sqlite_schema
WHERE type = 'table' AND name NOT LIKE 'sqlite_%'",
[],
|row| row.get(0),
)?;
ensure!(
table_count == 2,
"database contains missing or unexpected application tables"
);
let unexpected_objects: i64 = connection.query_row(
"SELECT count(*) FROM sqlite_schema
WHERE type IN ('view', 'trigger')
OR (type = 'index' AND name NOT LIKE 'sqlite_autoindex_%' AND name != 'vectors_search')",
[],
|row| row.get(0),
)?;
ensure!(
unexpected_objects == 0,
"database contains unexpected schema objects"
);
validate_schema_sql(connection, "table", "_migrations", MIGRATION_LEDGER_SQL)?;
validate_schema_sql(connection, "table", "vectors", VECTORS_TABLE_SQL)?;
validate_schema_sql(connection, "index", "vectors_search", VECTORS_INDEX_SQL)?;
Ok(())
}
fn validate_schema_sql(
connection: &Connection,
object_type: &str,
name: &str,
expected: &str,
) -> Result<()> {
let actual: String = connection
.query_row(
"SELECT sql FROM sqlite_schema WHERE type = ?1 AND name = ?2",
params![object_type, name],
|row| row.get(0),
)
.with_context(|| format!("database is missing {object_type} {name}"))?;
let actual = normalize_sql(&actual);
let expected = normalize_sql(expected);
ensure!(
actual == expected,
"database schema drift for {object_type} {name}: expected {expected:?}, got {actual:?}"
);
Ok(())
}
fn normalize_sql(sql: &str) -> String {
sql.trim_end_matches(';')
.split_whitespace()
.collect::<Vec<_>>()
.join(" ")
.replace("CREATE TABLE IF NOT EXISTS", "CREATE TABLE")
}
fn is_better(left: &(String, f32), right: &(String, f32)) -> bool {
left.1.total_cmp(&right.1).is_gt() || (left.1.total_cmp(&right.1).is_eq() && left.0 < right.0)
}
fn normalize_embedding(embedding: &[f32]) -> Result<Vec<f32>> {
ensure!(
(1..=MAX_DIMENSIONS).contains(&embedding.len()),
"embedding dimensions must be between 1 and {MAX_DIMENSIONS}"
);
ensure!(
embedding.iter().all(|value| value.is_finite()),
"embedding values must be finite"
);
let norm = embedding
.iter()
.map(|value| value * value)
.sum::<f32>()
.sqrt();
ensure!(
norm.is_finite() && norm > f32::EPSILON,
"embedding norm must be nonzero"
);
Ok(embedding.iter().map(|value| value / norm).collect())
}
fn encode_embedding(embedding: &[f32]) -> Vec<u8> {
embedding
.iter()
.flat_map(|value| value.to_le_bytes())
.collect()
}
/// One shared bound for writes and searches: the number of vectors a
/// (tenant, model, dimension) scope may hold so an exhaustive scan stays
/// within `MAX_SEARCH_WORK` multiply-adds and `MAX_CANDIDATES` rows.
fn vector_capacity(dimensions: usize) -> usize {
(MAX_SEARCH_WORK / dimensions.max(1)).min(MAX_CANDIDATES)
}
fn validate_identifier(label: &str, value: &str) -> Result<()> {
ensure!(
(1..=128).contains(&value.len()),
"{label} must contain 1 to 128 bytes"
);
ensure!(
value
.bytes()
.all(|byte| byte.is_ascii_alphanumeric() || b"-_.:".contains(&byte)),
"{label} contains unsupported characters"
);
Ok(())
}
fn validate_secret(value: &str) -> Result<()> {
ensure!(
value.len() == 64 && value.bytes().all(|byte| byte.is_ascii_hexdigit()),
"secret must be 32 random bytes encoded as 64 hexadecimal characters"
);
Ok(())
}
fn decode_secret(value: &str) -> Result<Zeroizing<Vec<u8>>> {
validate_secret(value)?;
let mut decoded = Zeroizing::new(Vec::with_capacity(32));
for pair in value.as_bytes().chunks_exact(2) {
let high = hex_nibble(pair[0])?;
let low = hex_nibble(pair[1])?;
decoded.push((high << 4) | low);
}
Ok(decoded)
}
fn hex_nibble(value: u8) -> Result<u8> {
match value {
b'0'..=b'9' => Ok(value - b'0'),
b'a'..=b'f' => Ok(value - b'a' + 10),
b'A'..=b'F' => Ok(value - b'A' + 10),
_ => Err(anyhow!("secret contains a non-hexadecimal byte")),
}
}
fn prepare_database_path(path: &Path) -> Result<()> {
ensure!(
!path.as_os_str().is_empty(),
"database path cannot be empty"
);
if metadata_if_exists(path)?.is_some() {
let metadata = fs::symlink_metadata(path)?;
ensure!(
!metadata.file_type().is_symlink(),
"database cannot be a symlink"
);
ensure!(metadata.is_file(), "database path is not a file");
ensure_private_mode(path, &metadata)?;
ensure_single_link(path, "database", &metadata)?;
} else {
let parent = path
.parent()
.filter(|parent| !parent.as_os_str().is_empty());
if let Some(parent) = parent
&& !parent.exists()
{
create_private_dir(parent)?;
}
create_private_file(path)?;
}
ensure_secure_parent(path)?;
Ok(())
}
fn prepare_output_parent(path: &Path) -> Result<()> {
ensure!(!path.as_os_str().is_empty(), "output path cannot be empty");
if let Some(parent) = path
.parent()
.filter(|parent| !parent.as_os_str().is_empty())
&& !parent.exists()
{
create_private_dir(parent)?;
}
ensure_secure_parent(path)?;
Ok(())
}
fn open_lock(database: &Path) -> Result<File> {
let lock_path = sibling_path(database, ".lock")?;
if metadata_if_exists(&lock_path)?.is_some() {
let metadata = fs::symlink_metadata(&lock_path)?;
ensure!(
!metadata.file_type().is_symlink() && metadata.is_file(),
"database lock must be a regular file"
);
ensure_private_mode(&lock_path, &metadata)?;
#[cfg(unix)]
{
use std::os::unix::fs::MetadataExt;
ensure!(metadata.nlink() == 1, "database lock cannot be hard-linked");
}
crate::fsguard::open_nofollow(&lock_path, true)
} else {
open_new_private(&lock_path)
}
}
fn open_validated_existing_file(path: &Path, label: &str, write: bool) -> Result<File> {
let file = crate::fsguard::open_nofollow(path, write)?;
let metadata = file.metadata()?;
ensure!(metadata.len() > 0, "{label} is empty");
ensure_private_mode(path, &metadata)?;
ensure_single_link(path, label, &metadata)?;
crate::fsguard::ensure_path_matches_file(path, &file, label)?;
Ok(file)
}
fn ensure_single_link(path: &Path, label: &str, metadata: &fs::Metadata) -> Result<()> {
#[cfg(unix)]
{
use std::os::unix::fs::MetadataExt;
ensure!(
metadata.nlink() == 1,
"{label} {} cannot be hard-linked",
path.display()
);
}
Ok(())
}
fn ensure_files_differ(
left_path: &Path,
left: &File,
right_path: &Path,
right: &File,
message: &str,
) -> Result<()> {
let left = left.metadata()?;
let right = right.metadata()?;
#[cfg(unix)]
ensure!(
{
use std::os::unix::fs::MetadataExt;
left.dev() != right.dev() || left.ino() != right.ino()
},
"{message}"
);
#[cfg(not(unix))]
ensure!(
fs::canonicalize(left_path)? != fs::canonicalize(right_path)?,
"{message}"
);
#[cfg(unix)]
let _ = (left_path, right_path);
Ok(())
}
fn copy_file_contents(source: &mut File, destination: &mut File) -> Result<()> {
source.seek(SeekFrom::Start(0))?;
destination.set_len(0)?;
destination.seek(SeekFrom::Start(0))?;
std::io::copy(source, destination)?;
destination.sync_all()?;
Ok(())
}
fn remove_file_if_matches(path: &Path, file: &File) {
if crate::fsguard::ensure_path_matches_file(path, file, "temporary file").is_ok() {
let _ = fs::remove_file(path);
}
}
fn create_private_dir(path: &Path) -> Result<()> {
fs::create_dir_all(path).with_context(|| format!("create {}", path.display()))?;
#[cfg(unix)]
{
use std::os::unix::fs::PermissionsExt;
fs::set_permissions(path, fs::Permissions::from_mode(0o700))?;
}
Ok(())
}
fn create_private_file(path: &Path) -> Result<()> {
drop(open_new_private(path)?);
Ok(())
}
fn open_new_private(path: &Path) -> Result<File> {
let mut options = OpenOptions::new();
options.read(true).write(true).create_new(true);
#[cfg(unix)]
{
use std::os::unix::fs::OpenOptionsExt;
options.mode(0o600);
}
options
.open(path)
.with_context(|| format!("create private file {}", path.display()))
}
fn ensure_private_mode(path: &Path, metadata: &fs::Metadata) -> Result<()> {
#[cfg(unix)]
{
use std::os::unix::fs::PermissionsExt;
let mode = metadata.permissions().mode();
ensure!(
mode & 0o077 == 0,
"secret file {} must not be accessible by group or others",
path.display()
);
}
Ok(())
}
fn validate_existing_file(path: &Path, label: &str) -> Result<()> {
let metadata = fs::symlink_metadata(path)
.with_context(|| format!("inspect {label} {}", path.display()))?;
ensure!(
!metadata.file_type().is_symlink(),
"{label} cannot be a symlink"
);
ensure!(metadata.is_file(), "{label} is not a regular file");
ensure!(metadata.len() > 0, "{label} is empty");
ensure_private_mode(path, &metadata)?;
#[cfg(unix)]
{
use std::os::unix::fs::MetadataExt;
ensure!(metadata.nlink() == 1, "{label} cannot be hard-linked");
}
Ok(())
}
fn metadata_if_exists(path: &Path) -> Result<Option<fs::Metadata>> {
match fs::symlink_metadata(path) {
Ok(metadata) => Ok(Some(metadata)),
Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(None),
Err(error) => Err(error).with_context(|| format!("inspect {}", path.display())),
}
}
fn ensure_secure_parent(path: &Path) -> Result<()> {
let parent = path
.parent()
.filter(|parent| !parent.as_os_str().is_empty())
.unwrap_or_else(|| Path::new("."));
let metadata = fs::symlink_metadata(parent)
.with_context(|| format!("inspect parent directory {}", parent.display()))?;
ensure!(
!metadata.file_type().is_symlink() && metadata.is_dir(),
"data parent must be a real directory: {}",
parent.display()
);
#[cfg(unix)]
{
use std::os::unix::fs::PermissionsExt;
ensure!(
metadata.permissions().mode() & 0o022 == 0,
"data parent cannot be writable by group or others: {}",
parent.display()
);
}
Ok(())
}
fn sibling_path(path: &Path, suffix: &str) -> Result<PathBuf> {
let name = path
.file_name()
.ok_or_else(|| anyhow!("path has no file name: {}", path.display()))?;
let mut sibling = name.to_os_string();
sibling.push(suffix);
Ok(path.with_file_name(sibling))
}
fn ensure_no_sidecars(path: &Path) -> Result<()> {
for suffix in ["-wal", "-shm"] {
let sidecar = sibling_path(path, suffix)?;
ensure!(
fs::symlink_metadata(&sidecar).is_err(),
"database sidecar remains after leaving WAL mode: {}",
sidecar.display()
);
}
Ok(())
}
fn sync_parent(path: &Path) -> Result<()> {
#[cfg(unix)]
{
let parent = path
.parent()
.filter(|parent| !parent.as_os_str().is_empty())
.unwrap_or_else(|| Path::new("."));
File::open(parent)?.sync_all()?;
}
Ok(())
}
fn unix_timestamp() -> Result<u64> {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.context("system clock is before Unix epoch")
.map(|duration| duration.as_secs())
}
#[cfg(test)]
mod tests {
use super::*;
const KEY: &str = "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef";
#[test]
fn migrates_and_searches_vectors() -> Result<()> {
let database = Database::memory(KEY)?;
assert_eq!(database.migrate()?, 1);
assert_eq!(database.migrate()?, 0);
database.upsert_vector("tenant", "model", "north", "North", &[1.0, 0.0])?;
database.upsert_vector("tenant", "model", "east", "East", &[0.0, 1.0])?;
let matches = database.search_vectors("tenant", "model", &[0.9, 0.1], 1)?;
assert_eq!(matches.len(), 1);
assert_eq!(matches[0].id, "north");
database.validate_schema()
}
#[test]
fn readiness_probes_the_open_connection() -> Result<()> {
Database::memory(KEY)?.check_ready()
}
#[test]
fn vector_capacity_tracks_the_search_budget() {
assert_eq!(vector_capacity(1), 10_000);
assert_eq!(vector_capacity(1_000), 10_000);
assert_eq!(vector_capacity(1_024), 9_765);
assert_eq!(vector_capacity(4_096), 2_441);
}
#[test]
fn rejects_unsafe_vectors() -> Result<()> {
let database = Database::memory(KEY)?;
database.migrate()?;
assert!(
database
.upsert_vector("tenant", "model", "zero", "", &[0.0])
.is_err()
);
assert!(
database
.upsert_vector("tenant", "model", "nan", "", &[f32::NAN])
.is_err()
);
Ok(())
}
#[test]
fn vector_top_k_is_bounded_and_deterministic() -> Result<()> {
let database = Database::memory(KEY)?;
database.migrate()?;
database.upsert_vector("tenant", "model", "b", "B", &[1.0, 0.0])?;
database.upsert_vector("tenant", "model", "a", "A", &[1.0, 0.0])?;
database.upsert_vector("tenant", "model", "south", "South", &[-1.0, 0.0])?;
let matches = database.search_vectors("tenant", "model", &[1.0, 0.0], 2)?;
assert_eq!(
matches
.iter()
.map(|item| item.id.as_str())
.collect::<Vec<_>>(),
["a", "b"]
);
assert!(
database
.search_vectors("tenant", "model", &[1.0, 0.0], MAX_RESULTS + 1)
.is_err()
);
Ok(())
}
#[test]
fn encrypted_backup_round_trip() -> Result<()> {
let directory = tempfile::tempdir()?;
let live = directory.path().join("live.db");
let backup = directory.path().join("backup.db");
let database = Database::open_for_migration(&live, KEY)?;
database.migrate()?;
database.upsert_vector("tenant", "model", "one", "One", &[1.0, 2.0])?;
database.backup(&backup, KEY)?;
let restored = Database::open_existing(&backup, KEY)?;
let matches = restored.search_vectors("tenant", "model", &[1.0, 2.0], 1)?;
assert_eq!(matches.len(), 1);
assert_eq!(matches[0].id, "one");
Ok(())
}
#[test]
fn wrong_key_is_rejected() -> Result<()> {
let directory = tempfile::tempdir()?;
let path = directory.path().join("encrypted.db");
let database = Database::open_for_migration(&path, KEY)?;
database.migrate()?;
drop(database);
let wrong = "ffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff";
assert!(Database::open_existing(&path, wrong).is_err());
Ok(())
}
#[test]
fn key_hex_is_case_insensitive() -> Result<()> {
let directory = tempfile::tempdir()?;
let path = directory.path().join("encrypted.db");
let database = Database::open_for_migration(&path, KEY)?;
database.migrate()?;
drop(database);
Database::open_existing(&path, &KEY.to_uppercase())?;
Ok(())
}
#[test]
fn secret_equality_compares_decoded_bytes() -> Result<()> {
assert!(secrets_equal(KEY, &KEY.to_uppercase())?);
assert!(!secrets_equal(
KEY,
"ffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff"
)?);
assert!(secrets_equal(KEY, "not-hex").is_err());
Ok(())
}
#[cfg(unix)]
#[test]
fn migration_rejects_hard_linked_database() -> Result<()> {
let directory = tempfile::tempdir()?;
let original = directory.path().join("original.db");
let linked = directory.path().join("linked.db");
let database = Database::open_for_migration(&original, KEY)?;
database.migrate()?;
drop(database);
fs::hard_link(&original, &linked)?;
let error = Database::open_for_migration(&linked, KEY)
.err()
.context("migration accepted a hard-linked database")?;
assert!(error.to_string().contains("hard-linked"));
Ok(())
}
#[test]
fn schema_validation_rejects_unexpected_objects() -> Result<()> {
let database = Database::memory(KEY)?;
database.migrate()?;
database
.connection()?
.execute_batch("CREATE TABLE unexpected(value TEXT) STRICT")?;
assert!(database.validate_schema().is_err());
Ok(())
}
#[cfg(unix)]
#[test]
fn open_file_identity_detects_path_replacement() -> Result<()> {
let directory = tempfile::tempdir()?;
let path = directory.path().join("pinned.db");
let moved = directory.path().join("moved.db");
let mut pinned = open_new_private(&path)?;
pinned.write_all(b"pinned")?;
fs::rename(&path, &moved)?;
let mut replacement = open_new_private(&path)?;
replacement.write_all(b"replacement")?;
assert!(crate::fsguard::ensure_path_matches_file(&path, &pinned, "pinned file").is_err());
Ok(())
}
#[test]
fn restore_rejects_empty_and_unrelated_databases() -> Result<()> {
let directory = tempfile::tempdir()?;
let live = directory.path().join("live.db");
let empty = directory.path().join("empty.db");
let unrelated = directory.path().join("unrelated.db");
let missing = directory.path().join("missing.db");
let database = Database::open_for_migration(&live, KEY)?;
database.migrate()?;
database.upsert_vector("tenant", "model", "safe", "Safe", &[1.0])?;
drop(database);
fs::write(&empty, [])?;
assert!(restore_database(&live, &missing, KEY).is_err());
assert!(!missing.exists());
assert!(restore_database(&live, &empty, KEY).is_err());
let connection = open_connection(&unrelated, KEY, false, true)?;
connection.execute_batch("CREATE TABLE unrelated(value TEXT) STRICT")?;
drop(connection);
assert!(restore_database(&live, &unrelated, KEY).is_err());
let database = Database::open_existing(&live, KEY)?;
assert_eq!(
database.search_vectors("tenant", "model", &[1.0], 1)?[0].id,
"safe"
);
Ok(())
}
#[cfg(unix)]
#[test]
fn non_utf8_backup_path_leaves_no_orphan() -> Result<()> {
use std::os::unix::ffi::{OsStrExt, OsStringExt};
let directory = tempfile::tempdir()?;
let live = directory.path().join("live.db");
let destination = directory
.path()
.join(std::ffi::OsString::from_vec(b"bad-\xff.db".to_vec()));
let database = Database::open_for_migration(&live, KEY)?;
database.migrate()?;
assert!(database.backup(&destination, KEY).is_err());
assert!(!destination.exists());
for entry in fs::read_dir(directory.path())? {
assert!(!entry?.file_name().as_bytes().starts_with(b"bad-"));
}
Ok(())
}
#[test]
fn restore_round_trip_preserves_previous_database() -> Result<()> {
let directory = tempfile::tempdir()?;
let live = directory.path().join("live.db");
let backup = directory.path().join("backup.db");
let database = Database::open_for_migration(&live, KEY)?;
database.migrate()?;
database.upsert_vector("tenant", "model", "before", "Before", &[1.0])?;
database.backup(&backup, KEY)?;
database.upsert_vector("tenant", "model", "after", "After", &[1.0])?;
drop(database);
let previous = restore_database(&live, &backup, KEY)?;
assert!(previous.is_file());
let restored = Database::open_existing(&live, KEY)?;
let matches = restored.search_vectors("tenant", "model", &[1.0], 10)?;
assert_eq!(matches.len(), 1);
assert_eq!(matches[0].id, "before");
Ok(())
}
}