use super::checkpoint::{append_padding, PartBytes, StorageCheckpoint};
use super::{segmented, SqliteStorageAdapter, StorageError, StorageResult};
use nostr::hashes::{sha256, Hash};
use rusqlite::{types::ValueRef, Connection, DatabaseName, OptionalExtension};
use std::collections::{BTreeMap, BTreeSet};
type Namespace<'a> = (&'a str, &'a str, &'a str);
const CHUNK: usize = 64 * 1024;
const PAGE: usize = 4096;
fn invalid() -> StorageError {
StorageError::new("invalid record checkpoint")
}
fn err(error: impl std::fmt::Display + std::any::Any) -> StorageError {
SqliteStorageAdapter::map_err(error)
}
fn number(value: i64) -> StorageResult<usize> {
usize::try_from(value).map_err(|_| invalid())
}
fn schema(conn: &Connection, create: bool) -> StorageResult<bool> {
let (count, valid): (i64, i64) = conn.query_row(
"SELECT count(*),coalesce(sum((name IN ('ndr_kv_checkpoint','ndr_kv_checkpoint_nodes','ndr_kv_checkpoint_chunks') AND type='table') OR (name='ndr_kv_checkpoint_delete' AND type='trigger')),0)
FROM sqlite_schema WHERE name IN ('ndr_kv_checkpoint','ndr_kv_checkpoint_nodes','ndr_kv_checkpoint_chunks','ndr_kv_checkpoint_delete')", [], |r| Ok((r.get(0)?,r.get(1)?)),
).map_err(err)?;
if count == 4 && valid == 4 {
super::checkpoint_index::schema(conn, create)?;
return Ok(true);
}
if count != 0 {
return Err(invalid());
}
if !create {
super::checkpoint_index::schema(conn, false)?;
return Ok(false);
}
conn.execute_batch(
"CREATE TABLE ndr_kv_checkpoint(owner_pubkey_hex TEXT NOT NULL,device_pubkey_hex TEXT NOT NULL,key TEXT NOT NULL,
format_version INTEGER NOT NULL,byte_length INTEGER NOT NULL,node_count INTEGER NOT NULL,head BLOB,
PRIMARY KEY(owner_pubkey_hex,device_pubkey_hex,key));
CREATE TABLE ndr_kv_checkpoint_nodes(owner_pubkey_hex TEXT NOT NULL,device_pubkey_hex TEXT NOT NULL,key TEXT NOT NULL,
node_key BLOB NOT NULL,next_key BLOB,payload_key BLOB NOT NULL,byte_length INTEGER NOT NULL,
prefix TEXT NOT NULL,suffix TEXT NOT NULL,padding INTEGER NOT NULL,
PRIMARY KEY(owner_pubkey_hex,device_pubkey_hex,key,node_key));
CREATE TABLE ndr_kv_checkpoint_chunks(owner_pubkey_hex TEXT NOT NULL,device_pubkey_hex TEXT NOT NULL,key TEXT NOT NULL,
payload_key BLOB NOT NULL,chunk_index INTEGER NOT NULL,data BLOB NOT NULL,
PRIMARY KEY(owner_pubkey_hex,device_pubkey_hex,key,payload_key,chunk_index));
CREATE TRIGGER ndr_kv_checkpoint_delete AFTER DELETE ON ndr_kv BEGIN
DELETE FROM ndr_kv_checkpoint_chunks WHERE owner_pubkey_hex=OLD.owner_pubkey_hex AND device_pubkey_hex=OLD.device_pubkey_hex AND key=OLD.key;
DELETE FROM ndr_kv_checkpoint_nodes WHERE owner_pubkey_hex=OLD.owner_pubkey_hex AND device_pubkey_hex=OLD.device_pubkey_hex AND key=OLD.key;
DELETE FROM ndr_kv_checkpoint WHERE owner_pubkey_hex=OLD.owner_pubkey_hex AND device_pubkey_hex=OLD.device_pubkey_hex AND key=OLD.key;
END;"
).map_err(err)?;
super::checkpoint_index::schema(conn, create)?;
Ok(true)
}
#[derive(PartialEq, Eq)]
struct Meta {
bytes: usize,
count: usize,
head: Option<Vec<u8>>,
}
struct Node<T> {
rowid: i64,
next: Option<Vec<u8>>,
payload: Vec<u8>,
bytes: usize,
prefix_len: usize,
suffix_len: usize,
padding: usize,
descriptors: T,
}
#[derive(Clone, Copy, PartialEq, Eq)]
struct NodeView<'a> {
next: Option<&'a [u8]>,
payload: &'a [u8],
bytes: usize,
prefix: &'a str,
suffix: &'a str,
padding: usize,
}
impl NodeView<'_> {
fn stored_bytes(&self) -> usize {
self.next.map_or(0, <[u8]>::len)
+ self.payload.len()
+ self.prefix.len()
+ self.suffix.len()
+ 16
}
}
impl<T> Node<T> {
fn logical_bytes(&self) -> StorageResult<usize> {
self.bytes
.checked_add(self.prefix_len)
.and_then(|n| n.checked_add(self.suffix_len))
.and_then(|n| n.checked_add(self.padding))
.ok_or_else(invalid)
}
}
fn proposed_node(value: &StorageCheckpoint, index: usize) -> NodeView<'_> {
let part = &value.parts[index];
NodeView {
next: value.parts.get(index + 1).map(|p| p.key.as_slice()),
payload: &part.payload_key,
bytes: part.bytes.bytes().len(),
prefix: &part.prefix,
suffix: &part.suffix,
padding: part.padding,
}
}
fn owned_descriptors(_: &[u8], node: NodeView<'_>) -> (String, String) {
(node.prefix.to_owned(), node.suffix.to_owned())
}
fn metadata(conn: &Connection, ns: Namespace<'_>) -> StorageResult<Option<Meta>> {
if !schema(conn, false)? {
return Ok(None);
}
let row: Option<(i64,i64,i64,Option<Vec<u8>>)> = conn.query_row(
"SELECT format_version,byte_length,node_count,head FROM ndr_kv_checkpoint WHERE owner_pubkey_hex=?1 AND device_pubkey_hex=?2 AND key=?3", ns,
|r| Ok((r.get(0)?,r.get(1)?,r.get(2)?,r.get(3)?)),
).optional().map_err(err)?;
if let Some((version, bytes, count, head)) = row {
if version != 1 {
return Err(invalid());
}
if segmented::has_value_in_snapshot(conn, ns)? {
return Err(invalid());
}
let anchor: bool = conn.query_row("SELECT EXISTS(SELECT 1 FROM ndr_kv WHERE owner_pubkey_hex=?1 AND device_pubkey_hex=?2 AND key=?3 AND typeof(value)='text' AND length(CAST(value AS BLOB))=0)",ns,|r|r.get(0)).map_err(err)?;
if !anchor {
return Err(invalid());
}
segmented::account_io(24 + head.as_ref().map_or(0, Vec::len), 0);
return Ok(Some(Meta {
bytes: number(bytes)?,
count: number(count)?,
head,
}));
}
let orphan: bool = conn.query_row("SELECT EXISTS(SELECT 1 FROM ndr_kv_checkpoint_nodes WHERE owner_pubkey_hex=?1 AND device_pubkey_hex=?2 AND key=?3) OR EXISTS(SELECT 1 FROM ndr_kv_checkpoint_chunks WHERE owner_pubkey_hex=?1 AND device_pubkey_hex=?2 AND key=?3)",ns,|r|r.get(0)).map_err(err)?;
if orphan {
return Err(invalid());
}
Ok(None)
}
fn nodes<T>(
conn: &Connection,
ns: Namespace<'_>,
meta: &Meta,
mut descriptors: impl FnMut(&[u8], NodeView<'_>) -> T,
) -> StorageResult<BTreeMap<Vec<u8>, Node<T>>> {
let query = if super::checkpoint_index::schema(conn, false)? {
"SELECT node_key,next_key,payload_key,byte_length,prefix,suffix,padding,rowid FROM ndr_kv_checkpoint_nodes INDEXED BY ndr_kv_checkpoint_node_namespace WHERE owner_pubkey_hex=?1 AND device_pubkey_hex=?2 AND key=?3 ORDER BY rowid"
} else {
"SELECT node_key,next_key,payload_key,byte_length,prefix,suffix,padding,rowid FROM ndr_kv_checkpoint_nodes WHERE owner_pubkey_hex=?1 AND device_pubkey_hex=?2 AND key=?3"
};
let mut stmt = conn.prepare(query).map_err(err)?;
let mut rows = stmt.query(ns).map_err(err)?;
let mut result = BTreeMap::new();
while let Some(row) = rows.next().map_err(err)? {
let key: Vec<u8> = row.get(0).map_err(err)?;
let next: Option<Vec<u8>> = row.get(1).map_err(err)?;
let payload: Vec<u8> = row.get(2).map_err(err)?;
let view = NodeView {
next: next.as_deref(),
payload: &payload,
bytes: number(row.get(3).map_err(err)?)?,
prefix: row.get_ref(4).map_err(err)?.as_str().map_err(err)?,
suffix: row.get_ref(5).map_err(err)?.as_str().map_err(err)?,
padding: number(row.get(6).map_err(err)?)?,
};
if !matches!(view.payload.first(), Some(0 | 1))
|| (view.payload[0] == 1 && view.payload.len() != 33)
{
return Err(invalid());
}
segmented::account_io(key.len() + view.stored_bytes(), 0);
let node = Node {
rowid: row.get(7).map_err(err)?,
bytes: view.bytes,
prefix_len: view.prefix.len(),
suffix_len: view.suffix.len(),
padding: view.padding,
descriptors: descriptors(&key, view),
next,
payload,
};
if result.insert(key, node).is_some() {
return Err(invalid());
}
}
if result.len() != meta.count {
return Err(invalid());
}
let mut seen = BTreeSet::new();
let mut next = meta.head.as_ref();
let mut total = 0usize;
while let Some(key) = next {
let node = result.get(key).ok_or_else(invalid)?;
if !seen.insert(key) {
return Err(invalid());
}
total = total
.checked_add(node.logical_bytes()?)
.ok_or_else(invalid)?;
next = node.next.as_ref();
}
if seen.len() != result.len() || total != meta.bytes {
return Err(invalid());
}
Ok(result)
}
enum PayloadScan {
RawShapes,
IndexedShapes,
References,
}
struct PayloadCheck {
bytes: usize,
next_chunk: usize,
}
fn payloads<'a, T>(
conn: &Connection,
ns: Namespace<'_>,
nodes: &'a BTreeMap<Vec<u8>, Node<T>>,
scan: PayloadScan,
) -> StorageResult<BTreeMap<&'a [u8], PayloadCheck>> {
let mut expected = BTreeMap::new();
for node in nodes.values() {
let payload = expected
.entry(node.payload.as_slice())
.or_insert(PayloadCheck {
bytes: node.bytes,
next_chunk: 0,
});
if payload.bytes != node.bytes {
return Err(invalid());
}
}
let query = match scan {
PayloadScan::IndexedShapes => "SELECT payload_key,chunk_index,byte_length,data_type FROM ndr_kv_checkpoint_chunk_shapes WHERE owner_pubkey_hex=?1 AND device_pubkey_hex=?2 AND key=?3 ORDER BY payload_key,chunk_index",
PayloadScan::RawShapes => "SELECT payload_key,chunk_index,length(data),typeof(data) FROM ndr_kv_checkpoint_chunks WHERE owner_pubkey_hex=?1 AND device_pubkey_hex=?2 AND key=?3 ORDER BY payload_key,chunk_index",
PayloadScan::References => "SELECT payload_key,chunk_index FROM ndr_kv_checkpoint_chunks WHERE owner_pubkey_hex=?1 AND device_pubkey_hex=?2 AND key=?3 ORDER BY payload_key,chunk_index",
};
let mut stmt = conn.prepare(query).map_err(err)?;
let mut rows = stmt.query(ns).map_err(err)?;
while let Some(row) = rows.next().map_err(err)? {
let ValueRef::Blob(key) = row.get_ref(0).map_err(err)? else {
return Err(invalid());
};
let index = number(row.get(1).map_err(err)?)?;
let payload = expected.get_mut(key).ok_or_else(invalid)?;
if index != payload.next_chunk || index >= payload.bytes.div_ceil(CHUNK) {
return Err(invalid());
}
if !matches!(scan, PayloadScan::References) {
let bytes = number(row.get(2).map_err(err)?)?;
let ValueRef::Text(kind) = row.get_ref(3).map_err(err)? else {
return Err(invalid());
};
if bytes != (payload.bytes - index * CHUNK).min(CHUNK) || kind != b"blob" {
return Err(invalid());
}
}
payload.next_chunk += 1;
let metadata_bytes = if matches!(scan, PayloadScan::References) {
8
} else {
20
};
segmented::account_io(key.len() + metadata_bytes, 0);
}
if expected
.values()
.any(|payload| payload.next_chunk != payload.bytes.div_ceil(CHUNK))
{
return Err(invalid());
}
Ok(expected)
}
pub(super) fn get_in_snapshot(
conn: &Connection,
ns: Namespace<'_>,
) -> StorageResult<Option<String>> {
let Some(meta) = metadata(conn, ns)? else {
return Ok(None);
};
let nodes = nodes(conn, ns, &meta, owned_descriptors)?;
payloads(conn, ns, &nodes, PayloadScan::References)?;
let mut output = String::new();
output.try_reserve_exact(meta.bytes).map_err(err)?;
let mut next = meta.head.as_ref();
let mut stmt = conn.prepare("SELECT chunk_index,data FROM ndr_kv_checkpoint_chunks WHERE owner_pubkey_hex=?1 AND device_pubkey_hex=?2 AND key=?3 AND payload_key=?4 ORDER BY chunk_index").map_err(err)?;
let mut materialized = BTreeMap::<&[u8], std::ops::Range<usize>>::new();
while let Some(key) = next {
let node = nodes.get(key).ok_or_else(invalid)?;
output.push_str(&node.descriptors.0);
if let Some(range) = materialized.get(node.payload.as_slice()) {
output.extend_from_within(range.clone());
} else {
let mut bytes = Vec::new();
bytes.try_reserve_exact(node.bytes).map_err(err)?;
let mut rows = stmt.query((ns.0, ns.1, ns.2, &node.payload)).map_err(err)?;
let mut next_chunk = 0;
while let Some(row) = rows.next().map_err(err)? {
let index = number(row.get(0).map_err(err)?)?;
let ValueRef::Blob(data) = row.get_ref(1).map_err(err)? else {
return Err(invalid());
};
if index != next_chunk
|| index >= node.bytes.div_ceil(CHUNK)
|| data.len() != (node.bytes - index * CHUNK).min(CHUNK)
{
return Err(invalid());
}
next_chunk += 1;
segmented::account_io(data.len(), 0);
bytes.extend_from_slice(data);
}
if next_chunk != node.bytes.div_ceil(CHUNK)
|| bytes.len() != node.bytes
|| (node.payload[0] == 1
&& sha256::Hash::hash(&bytes).to_byte_array().as_slice() != &node.payload[1..])
{
return Err(invalid());
}
let start = output.len();
output.push_str(std::str::from_utf8(&bytes).map_err(|_| invalid())?);
materialized.insert(node.payload.as_slice(), start..output.len());
}
output.push_str(&node.descriptors.1);
append_padding(&mut output, node.padding);
next = node.next.as_ref();
}
if output.len() != meta.bytes {
return Err(invalid());
}
Ok(Some(output))
}
pub(super) fn clear_in_snapshot(conn: &Connection, ns: Namespace<'_>) -> StorageResult<()> {
let Some(meta) = metadata(conn, ns)? else {
return Ok(());
};
let old = nodes(conn, ns, &meta, owned_descriptors)?;
payloads(conn, ns, &old, PayloadScan::RawShapes)?;
for table in [
"ndr_kv_checkpoint_chunks",
"ndr_kv_checkpoint_nodes",
"ndr_kv_checkpoint",
] {
conn.execute(
&format!(
"DELETE FROM {table} WHERE owner_pubkey_hex=?1 AND device_pubkey_hex=?2 AND key=?3"
),
ns,
)
.map_err(err)?;
}
Ok(())
}
fn put_payload(
conn: &Connection,
ns: Namespace<'_>,
key: &[u8],
bytes: &[u8],
previous: Option<usize>,
) -> StorageResult<()> {
let mut changes = Vec::<(usize, Option<i64>, Vec<usize>)>::new();
if previous.is_some() {
let mut stmt = conn.prepare("SELECT rowid,chunk_index,data FROM ndr_kv_checkpoint_chunks WHERE owner_pubkey_hex=?1 AND device_pubkey_hex=?2 AND key=?3 AND payload_key=?4 ORDER BY chunk_index").map_err(err)?;
let mut rows = stmt.query((ns.0, ns.1, ns.2, key)).map_err(err)?;
while let Some(row) = rows.next().map_err(err)? {
let index = number(row.get(1).map_err(err)?)?;
let ValueRef::Blob(old) = row.get_ref(2).map_err(err)? else {
return Err(invalid());
};
segmented::account_io(old.len(), 0);
let Some(new) = bytes.get(index * CHUNK..bytes.len().min((index + 1) * CHUNK)) else {
continue;
};
if new.is_empty() {
continue;
}
if new.len() != old.len() {
changes.push((index, None, Vec::new()));
} else {
let pages = old
.chunks(PAGE)
.zip(new.chunks(PAGE))
.enumerate()
.filter_map(|(i, (a, b))| (a != b).then_some(i))
.collect::<Vec<_>>();
if !pages.is_empty() {
changes.push((index, Some(row.get(0).map_err(err)?), pages));
}
}
}
}
for index in previous.unwrap_or(0).div_ceil(CHUNK)..bytes.len().div_ceil(CHUNK) {
changes.push((index, None, Vec::new()));
}
for (index, rowid, pages) in changes {
let value = &bytes[index * CHUNK..bytes.len().min((index + 1) * CHUNK)];
if let Some(rowid) = rowid {
let mut blob = conn
.blob_open(
DatabaseName::Main,
"ndr_kv_checkpoint_chunks",
"data",
rowid,
false,
)
.map_err(err)?;
for page in pages {
let value = &value[page * PAGE..value.len().min((page + 1) * PAGE)];
blob.write_at(value, page * PAGE).map_err(err)?;
segmented::account_io(0, value.len());
}
blob.close().map_err(err)?;
} else {
conn.execute("INSERT INTO ndr_kv_checkpoint_chunks(owner_pubkey_hex,device_pubkey_hex,key,payload_key,chunk_index,data) VALUES(?1,?2,?3,?4,?5,?6) ON CONFLICT(owner_pubkey_hex,device_pubkey_hex,key,payload_key,chunk_index) DO UPDATE SET data=excluded.data",(ns.0,ns.1,ns.2,key,index as i64,value)).map_err(err)?;
segmented::account_io(0, value.len());
}
}
if previous.is_some_and(|old| old > bytes.len()) {
conn.execute("DELETE FROM ndr_kv_checkpoint_chunks WHERE owner_pubkey_hex=?1 AND device_pubkey_hex=?2 AND key=?3 AND payload_key=?4 AND chunk_index>=?5",(ns.0,ns.1,ns.2,key,bytes.len().div_ceil(CHUNK) as i64)).map_err(err)?;
}
Ok(())
}
pub(super) fn put(
conn: &mut Connection,
ns: Namespace<'_>,
value: &StorageCheckpoint,
) -> StorageResult<()> {
value.validate()?;
let tx = conn.savepoint().map_err(err)?;
schema(&tx, true)?;
let previous = metadata(&tx, ns)?;
let mut current = BTreeMap::new();
for (index, part) in value.parts.iter().enumerate() {
if current.insert(part.key.as_slice(), index).is_some() {
return Err(invalid());
}
}
let mut changed = vec![true; value.parts.len()];
let old = previous
.as_ref()
.map(|meta| {
nodes(&tx, ns, meta, |key, stored| {
if let Some(&index) = current.get(key) {
changed[index] = stored != proposed_node(value, index);
}
})
})
.transpose()?
.unwrap_or_default();
let old_payloads = payloads(&tx, ns, &old, PayloadScan::IndexedShapes)?;
let mut update_node = tx.prepare("UPDATE ndr_kv_checkpoint_nodes SET next_key=?1,payload_key=?2,byte_length=?3,prefix=?4,suffix=?5,padding=?6 WHERE rowid=?7").map_err(err)?;
let mut insert_node = tx.prepare("INSERT INTO ndr_kv_checkpoint_nodes(owner_pubkey_hex,device_pubkey_hex,key,node_key,next_key,payload_key,byte_length,prefix,suffix,padding) VALUES(?1,?2,?3,?4,?5,?6,?7,?8,?9,?10)").map_err(err)?;
let mut current_payloads = BTreeSet::new();
for (index, part) in value.parts.iter().enumerate() {
let node = proposed_node(value, index);
if current_payloads.insert(part.payload_key.as_slice()) {
let old_bytes = old_payloads
.get(part.payload_key.as_slice())
.map(|payload| payload.bytes);
if matches!(&part.bytes, PartBytes::Inline(_)) || old_bytes.is_none() {
put_payload(&tx, ns, &part.payload_key, part.bytes.bytes(), old_bytes)?;
} else if old_bytes != Some(node.bytes) {
return Err(invalid());
}
}
if changed[index] {
let affected = if let Some(previous) = old.get(part.key.as_slice()) {
update_node
.execute(rusqlite::params![
node.next,
node.payload,
node.bytes as i64,
node.prefix,
node.suffix,
node.padding as i64,
previous.rowid
])
.map_err(err)?
} else {
insert_node
.execute(rusqlite::params![
ns.0,
ns.1,
ns.2,
part.key,
node.next,
node.payload,
node.bytes as i64,
node.prefix,
node.suffix,
node.padding as i64
])
.map_err(err)?
};
if affected != 1 {
return Err(invalid());
}
segmented::account_io(0, part.key.len() + node.stored_bytes());
}
}
for key in old
.keys()
.filter(|key| !current.contains_key(key.as_slice()))
{
tx.execute("DELETE FROM ndr_kv_checkpoint_nodes WHERE owner_pubkey_hex=?1 AND device_pubkey_hex=?2 AND key=?3 AND node_key=?4",(ns.0,ns.1,ns.2,key)).map_err(err)?;
}
for key in old_payloads
.keys()
.copied()
.filter(|key| !current_payloads.contains(*key))
{
tx.execute("DELETE FROM ndr_kv_checkpoint_chunks WHERE owner_pubkey_hex=?1 AND device_pubkey_hex=?2 AND key=?3 AND payload_key=?4",(ns.0,ns.1,ns.2,key)).map_err(err)?;
}
let meta = Meta {
bytes: value.len(),
count: value.parts.len(),
head: value.parts.first().map(|p| p.key.clone()),
};
if previous.is_none() {
segmented::clear_in_snapshot(&tx, ns)?;
tx.execute("INSERT INTO ndr_kv(owner_pubkey_hex,device_pubkey_hex,key,value) VALUES(?1,?2,?3,'') ON CONFLICT(owner_pubkey_hex,device_pubkey_hex,key) DO UPDATE SET value=''",ns).map_err(err)?;
}
if previous.as_ref() != Some(&meta) {
tx.execute("INSERT INTO ndr_kv_checkpoint(owner_pubkey_hex,device_pubkey_hex,key,format_version,byte_length,node_count,head) VALUES(?1,?2,?3,1,?4,?5,?6) ON CONFLICT(owner_pubkey_hex,device_pubkey_hex,key) DO UPDATE SET byte_length=excluded.byte_length,node_count=excluded.node_count,head=excluded.head",(ns.0,ns.1,ns.2,meta.bytes as i64,meta.count as i64,&meta.head)).map_err(err)?;
segmented::account_io(0, 24 + meta.head.as_ref().map_or(0, Vec::len));
}
drop(update_node);
drop(insert_node);
tx.commit().map_err(err)
}