use std::path::Path;
use base64::Engine;
use rusqlite::{OptionalExtension, TransactionBehavior, params};
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
use tokio_rusqlite::Connection;
use crate::broker::now;
const TOKEN_SECRET_BYTES: usize = 32;
pub const DATABASE: &str = "review-store.sqlite";
const SCHEMA_VERSION: u32 = 3;
pub const MAX_ENTRIES: i64 = 100_000;
pub const MAX_BYTES: i64 = 4 * 1024 * 1024 * 1024;
#[derive(
Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize, clap::ValueEnum,
)]
#[serde(rename_all = "lowercase")]
pub enum Scope {
Read,
Publish,
Human,
}
impl Scope {
pub fn name(self) -> &'static str {
match self {
Self::Read => "read",
Self::Publish => "publish",
Self::Human => "human",
}
}
}
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct Token {
pub name: crate::types::TokenName,
pub scopes: Vec<Scope>,
pub created_at: crate::types::Timestamp,
pub revoked_at: Option<crate::types::Timestamp>,
}
#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct Revocation {
pub name: crate::types::TokenName,
pub revoked_at: crate::types::Timestamp,
pub purged_entries: usize,
}
#[derive(Debug, thiserror::Error)]
enum Error {
#[error(transparent)]
Sql(#[from] rusqlite::Error),
#[error(transparent)]
Json(#[from] serde_json::Error),
#[error("{0}")]
Invalid(String),
}
#[derive(Clone)]
pub struct Store {
connection: Connection,
}
fn digest(token: &str) -> crate::types::Sha256Digest {
Sha256::digest(token.as_bytes())
.iter()
.map(|byte| format!("{byte:02x}"))
.collect::<String>()
.parse()
.expect("SHA-256 token digest")
}
fn reuse_keys(db: &rusqlite::Transaction<'_>) -> Result<(), Error> {
db.execute_batch("DROP TABLE IF EXISTS entries;")?;
Ok(())
}
impl Store {
pub async fn open(state: &Path) -> Result<Self, String> {
let state = crate::store::state_dir(Some(state))?;
crate::platform::create_private_dir_all(&state).map_err(|e| e.to_string())?;
for suffix in ["", "-wal", "-shm"] {
let path = state.join(format!("{DATABASE}{suffix}"));
match crate::platform::path_kind(&path) {
Ok(crate::platform::FileKind::File) => {}
Err(error) if error.kind() == std::io::ErrorKind::NotFound => {}
Ok(_) => {
return Err(format!(
"Review store files must be regular files: {}",
crate::platform::path_text(&path)
));
}
Err(error) => return Err(error.to_string()),
}
}
let connection = Connection::open(state.join(DATABASE))
.await
.map_err(|e| e.to_string())?;
connection
.call(|db| -> Result<(), Error> {
db.busy_timeout(crate::store::SQLITE_BUSY_TIMEOUT)?;
let version: u32 = db.pragma_query_value(None, "user_version", |row| row.get(0))?;
if version > SCHEMA_VERSION {
return Err(Error::Invalid(format!(
"Unsupported review store schema version: {version}"
)));
}
db.pragma_update(None, "journal_mode", "WAL")?;
let transaction = db.transaction_with_behavior(TransactionBehavior::Immediate)?;
if matches!(
transaction
.pragma_query_value(None, "user_version", |row| row.get::<_, u32>(0))?,
1 | 2
) {
reuse_keys(&transaction)?;
}
transaction.execute_batch(
"CREATE TABLE IF NOT EXISTS tokens(name TEXT PRIMARY KEY, hash TEXT NOT NULL UNIQUE, scopes TEXT NOT NULL, created_at TEXT NOT NULL, revoked_at TEXT);
CREATE TABLE IF NOT EXISTS entries(key TEXT NOT NULL, execution_id TEXT NOT NULL, publisher TEXT NOT NULL, completed_at TEXT NOT NULL, bytes INTEGER NOT NULL, last_used TEXT NOT NULL, data TEXT NOT NULL, UNIQUE(key, execution_id));
CREATE INDEX IF NOT EXISTS entries_key ON entries(key, completed_at);
CREATE INDEX IF NOT EXISTS entries_lru ON entries(last_used, key);
CREATE INDEX IF NOT EXISTS entries_publisher ON entries(publisher);",
)?;
transaction.pragma_update(None, "user_version", SCHEMA_VERSION)?;
transaction.commit()?;
Ok(())
})
.await
.map_err(|e| e.to_string())?;
Ok(Self { connection })
}
pub async fn add_token(
&self,
name: &crate::types::TokenName,
scopes: &[Scope],
) -> Result<String, String> {
let mut scopes = scopes.to_vec();
scopes.sort();
scopes.dedup();
if scopes.is_empty() {
return Err("A token needs at least one scope.".into());
}
let mut secret = [0; TOKEN_SECRET_BYTES];
getrandom::fill(&mut secret).map_err(|_| "Cannot obtain secure randomness.".to_owned())?;
let token = format!(
"azt_{}",
base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(secret)
);
let (name, hash, scopes) = (
name.clone(),
digest(&token),
serde_json::to_string(&scopes).map_err(|e| e.to_string())?,
);
self.connection
.call(move |db| -> Result<(), Error> {
let inserted = db.execute(
"INSERT INTO tokens(name,hash,scopes,created_at) VALUES (?,?,?,?) ON CONFLICT(name) DO NOTHING",
params![name, hash, scopes, now()],
)?;
if inserted == 0 {
return Err(Error::Invalid(format!(
"Token name {name} is already used; names of revoked tokens are kept."
)));
}
Ok(())
})
.await
.map_err(|e| e.to_string())?;
Ok(token)
}
pub async fn tokens(&self) -> Result<Vec<Token>, String> {
self.connection
.call(|db| -> Result<_, Error> {
let mut statement = db.prepare(
"SELECT name,scopes,created_at,revoked_at FROM tokens ORDER BY name",
)?;
let rows = statement
.query_map([], |row| {
Ok((
row.get::<_, crate::types::TokenName>(0)?,
row.get::<_, String>(1)?,
row.get(2)?,
row.get(3)?,
))
})?
.map(|row| {
let (name, scopes, created_at, revoked_at) = row?;
Ok(Token {
name,
scopes: serde_json::from_str(&scopes)?,
created_at,
revoked_at,
})
})
.collect::<Result<_, Error>>()?;
Ok(rows)
})
.await
.map_err(|e| e.to_string())
}
pub async fn revoke(
&self,
name: &crate::types::TokenName,
purge: bool,
) -> Result<Revocation, String> {
let name = name.clone();
self.connection
.call(move |db| -> Result<_, Error> {
let transaction = db.transaction_with_behavior(TransactionBehavior::Immediate)?;
transaction.execute(
"UPDATE tokens SET revoked_at=? WHERE name=? AND revoked_at IS NULL",
params![now(), name],
)?;
let revoked_at: crate::types::Timestamp = transaction
.query_row(
"SELECT revoked_at FROM tokens WHERE name=?",
[&name],
|row| row.get(0),
)
.optional()?
.ok_or_else(|| Error::Invalid(format!("Unknown token {name}.")))?;
let purged_entries = if purge {
transaction.execute("DELETE FROM entries WHERE publisher=?", [&name])?
} else {
0
};
transaction.commit()?;
Ok(Revocation {
name,
revoked_at,
purged_entries,
})
})
.await
.map_err(|e| e.to_string())
}
pub async fn remove(&self, key: &crate::types::ReuseKey) -> Result<bool, String> {
let key = key.clone();
self.connection
.call(move |db| -> Result<_, Error> {
Ok(db.execute("DELETE FROM entries WHERE key=?", [key])? != 0)
})
.await
.map_err(|e| e.to_string())
}
pub(super) async fn authenticate(
&self,
token: &str,
) -> Result<Option<(crate::types::TokenName, Vec<Scope>)>, String> {
let hash = digest(token);
self.connection
.call(move |db| -> Result<_, Error> {
let row: Option<(crate::types::TokenName, String)> = db
.query_row(
"SELECT name,scopes FROM tokens WHERE hash=? AND revoked_at IS NULL",
[hash],
|row| Ok((row.get(0)?, row.get(1)?)),
)
.optional()?;
row.map(|(name, scopes)| Ok((name, serde_json::from_str(&scopes)?)))
.transpose()
})
.await
.map_err(|e| e.to_string())
}
pub(super) async fn lookup(
&self,
keys: Vec<crate::types::ReuseKey>,
) -> Result<Vec<String>, String> {
self.connection
.call(move |db| -> Result<_, Error> {
let transaction = db.transaction_with_behavior(TransactionBehavior::Immediate)?;
let used = now();
let mut found = Vec::new();
for key in keys {
let data: Option<String> = transaction
.query_row(
"UPDATE entries SET last_used=? WHERE rowid=(SELECT rowid FROM entries WHERE key=? ORDER BY completed_at DESC, rowid DESC LIMIT 1) RETURNING data",
params![used, key],
|row| row.get(0),
)
.optional()?;
found.extend(data);
}
transaction.commit()?;
Ok(found)
})
.await
.map_err(|e| e.to_string())
}
pub(super) async fn insert(
&self,
key: &crate::types::ReuseKey,
execution_id: &crate::types::ExecutionId,
publisher: &crate::types::TokenName,
completed_at: crate::types::Timestamp,
data: String,
) -> Result<bool, String> {
let record = (
key.clone(),
execution_id.clone(),
publisher.clone(),
completed_at,
);
self.connection
.call(move |db| -> Result<_, Error> {
let (key, execution_id, publisher, completed_at) = record;
let transaction = db.transaction_with_behavior(TransactionBehavior::Immediate)?;
let used = now();
let created = transaction.execute(
"INSERT INTO entries(key,execution_id,publisher,completed_at,bytes,last_used,data) VALUES (?,?,?,?,?,?,?) ON CONFLICT DO NOTHING",
params![
key,
execution_id,
publisher,
completed_at,
data.len() as i64,
used,
data
],
)? == 1;
loop {
let (count, bytes): (i64, i64) = transaction.query_row(
"SELECT count(*),coalesce(sum(bytes),0) FROM entries",
[],
|row| Ok((row.get(0)?, row.get(1)?)),
)?;
if count <= MAX_ENTRIES && bytes <= MAX_BYTES {
break;
}
if transaction.execute(
"DELETE FROM entries WHERE rowid=(SELECT rowid FROM entries WHERE NOT (key=?1 AND execution_id=?2) ORDER BY last_used,key,completed_at LIMIT 1)",
params![key, execution_id],
)? == 0
{
break;
}
}
transaction.commit()?;
Ok(created)
})
.await
.map_err(|e| e.to_string())
}
}