1use crate::*;
2
3cfg_if! {
4 if #[cfg(all(target_arch = "wasm32", target_os = "unknown"))] {
5 use keyvaluedb_web::*;
6 use keyvaluedb::*;
7 } else {
8 use keyvaluedb_sqlite::*;
9 use keyvaluedb::*;
10 }
11}
12
13impl_veilid_log_facility!("tstore");
14
15#[must_use]
16#[derive(Debug)]
17struct CryptInfo {
18 secret: SharedSecret,
19}
20impl CryptInfo {
21 pub fn new(secret: SharedSecret) -> Self {
22 Self { secret }
23 }
24}
25
26#[must_use]
28pub(super) struct TableDBUnlockedInner {
29 registry: VeilidComponentRegistry,
30 table: String,
31 database: Database,
32 commit_lock: AsyncMutex<()>,
34 encrypt_info: Option<CryptInfo>,
36 decrypt_info: Option<CryptInfo>,
37}
38
39impl fmt::Debug for TableDBUnlockedInner {
40 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
41 write!(f, "TableDBUnlockedInner(table={})", self.table)
42 }
43}
44
45#[derive(Debug, Clone)]
47#[must_use]
48pub struct TableDB {
49 opened_column_count: u32,
50 unlocked_inner: Arc<TableDBUnlockedInner>,
51}
52
53impl VeilidComponentRegistryAccessor for TableDB {
54 fn registry(&self) -> VeilidComponentRegistry {
55 self.unlocked_inner.registry.clone()
56 }
57}
58
59impl TableDB {
60 pub(super) fn new(
61 table: String,
62 registry: VeilidComponentRegistry,
63 database: Database,
64 encryption_key: Option<SharedSecret>,
65 decryption_key: Option<SharedSecret>,
66 opened_column_count: u32,
67 ) -> Self {
68 let encrypt_info = encryption_key.map(CryptInfo::new);
69 let decrypt_info = decryption_key.map(CryptInfo::new);
70
71 let total_columns = database.num_columns().unwrap_or_log();
72
73 Self {
74 opened_column_count: if opened_column_count == 0 {
75 total_columns
76 } else {
77 opened_column_count
78 },
79 unlocked_inner: Arc::new(TableDBUnlockedInner {
80 registry,
81 table,
82 database,
83 commit_lock: AsyncMutex::new(()),
84 encrypt_info,
85 decrypt_info,
86 }),
87 }
88 }
89
90 pub(super) fn new_from_unlocked_inner(
91 unlocked_inner: Arc<TableDBUnlockedInner>,
92 opened_column_count: u32,
93 ) -> Self {
94 let db = &unlocked_inner.database;
95 let total_columns = db.num_columns().unwrap_or_log();
96 Self {
97 opened_column_count: if opened_column_count == 0 {
98 total_columns
99 } else {
100 opened_column_count
101 },
102 unlocked_inner,
103 }
104 }
105
106 pub(super) fn unlocked_inner(&self) -> Arc<TableDBUnlockedInner> {
107 self.unlocked_inner.clone()
108 }
109
110 #[must_use]
112 pub fn table_name(&self) -> String {
113 self.unlocked_inner.table.clone()
114 }
115
116 #[cfg_attr(
118 feature = "instrument",
119 instrument(level = "trace", target = "tstore", skip_all)
120 )]
121 #[must_use]
122 pub fn io_stats(&self, kind: IoStatsKind) -> IoStats {
123 self.unlocked_inner.database.io_stats(kind)
124 }
125
126 pub async fn cleanup(&self) -> VeilidAPIResult<()> {
132 self.unlocked_inner
133 .database
134 .cleanup()
135 .measure_debug(
136 TimestampDuration::new_secs(1),
137 veilid_log_dbg!(self, "TableDB::cleanup {}", self.table_name()),
138 )
139 .await
140 .map_err(VeilidAPIError::internal)
141 }
142
143 #[cfg_attr(
148 feature = "instrument",
149 instrument(level = "trace", target = "tstore", skip_all)
150 )]
151 pub fn get_column_count(&self) -> VeilidAPIResult<u32> {
152 let db = &self.unlocked_inner.database;
153 db.num_columns().map_err(VeilidAPIError::from)
154 }
155
156 pub fn estimate_storage_size(
162 &self,
163 _col: u32,
164 key: &[u8],
165 value: &[u8],
166 ) -> VeilidAPIResult<u64> {
167 let size =
168 1 +
170 1 +
172 key.len() * 2 +
174 4 +
176 value.len() +
178 4 +
180 4;
183 size.try_into().map_err(VeilidAPIError::internal)
184 }
185
186 pub fn estimate_storage_size_json<T>(
190 &self,
191 col: u32,
192 key: &[u8],
193 value: &T,
194 ) -> VeilidAPIResult<u64>
195 where
196 T: serde::Serialize,
197 {
198 let value_json = serde_json::to_vec(value).map_err(VeilidAPIError::internal)?;
199 self.estimate_storage_size(col, key, &value_json)
200 }
201
202 #[cfg_attr(
209 feature = "instrument",
210 instrument(level = "trace", target = "tstore", skip_all)
211 )]
212 pub(in crate::table_store) async fn maybe_encrypt(
213 &self,
214 data: &[u8],
215 keyed_nonce: bool,
216 ) -> Bytes {
217 let Some(ei) = &self.unlocked_inner.encrypt_info else {
218 return Bytes::copy_from_slice(data);
219 };
220
221 let crypto = self.crypto();
222 let vcrypto = crypto.get_async(ei.secret.kind()).unwrap_or_log();
223 let mut out = BytesMut::zeroed(vcrypto.nonce_length() + data.len());
224
225 if keyed_nonce {
226 let mut noncedata = BytesMut::with_capacity(data.len() + ei.secret.ref_value().len());
228 noncedata.extend_from_slice(data);
229 noncedata.extend_from_slice(ei.secret.ref_value());
230 let noncehash = vcrypto.generate_hash(noncedata.freeze()).await.value();
231 out.as_mut()[0..vcrypto.nonce_length()]
233 .copy_from_slice(&noncehash.as_ref()[0..vcrypto.nonce_length()]);
234 } else {
235 random_bytes(&mut out[0..vcrypto.nonce_length()]);
237 }
238 let nonce = Nonce::new(&out[0..vcrypto.nonce_length()]);
239
240 let out = vcrypto
241 .crypt_b2b_no_auth(
242 Bytes::copy_from_slice(data),
243 out,
244 vcrypto.nonce_length(),
245 &nonce,
246 &ei.secret,
247 )
248 .await
249 .unwrap_or_log();
250
251 out.freeze()
252 }
253
254 #[cfg_attr(
256 feature = "instrument",
257 instrument(level = "trace", target = "tstore", skip_all)
258 )]
259 pub(in crate::table_store) async fn maybe_decrypt(
260 &self,
261 data: &[u8],
262 ) -> std::io::Result<Bytes> {
263 let Some(di) = &self.unlocked_inner.decrypt_info else {
264 return Ok(Bytes::copy_from_slice(data));
265 };
266
267 let crypto = self.crypto();
268 let vcrypto = crypto.get_async(di.secret.kind()).unwrap_or_log();
269 if data.len() < vcrypto.nonce_length() {
270 veilid_log!(self error "maybe_decrypt: data too short for nonce: {} < {}", data.len(), vcrypto.nonce_length());
271 return Err(std::io::Error::other("data too short for nonce"));
272 }
273 if data.len() == vcrypto.nonce_length() {
274 return Ok(Bytes::new());
275 }
276
277 let out = BytesMut::zeroed(data.len() - vcrypto.nonce_length());
278 let mut data = Bytes::copy_from_slice(data);
279 let data_start = data.split_to(vcrypto.nonce_length());
280
281 let out = vcrypto
282 .crypt_b2b_no_auth(data, out, 0, &Nonce::new(data_start.as_ref()), &di.secret)
283 .await
284 .unwrap_or_log();
285
286 Ok(out.freeze())
287 }
288
289 #[cfg_attr(
295 feature = "instrument",
296 instrument(level = "trace", target = "tstore", skip_all)
297 )]
298 pub async fn get_keys(&self, col: u32) -> VeilidAPIResult<Vec<Vec<u8>>> {
299 if col >= self.opened_column_count {
300 apibail_generic!(
301 "Column exceeds opened column count {} >= {}",
302 col,
303 self.opened_column_count
304 );
305 }
306 let db = self.unlocked_inner.database.clone();
307 let out = Vec::new();
308 let (mut out, _) = db
309 .iter_keys(col, None, out, |out, ekey| {
310 out.push(ekey.clone());
312 Ok(Option::<()>::None)
313 })
314 .await
315 .map_err(VeilidAPIError::from)?;
316
317 #[cfg(feature = "verbose-tracing")]
318 veilid_log!(self debug "TableDB::get_keys({}) col={}: read {} raw keys", self.unlocked_inner.table, col, out.len());
319
320 let max_value_size = self.config().table_store.max_value_size_mb as usize * 1024 * 1024;
321 for (idx, k) in out.iter_mut().enumerate() {
322 let raw_len = k.len();
323 let decrypted = self.maybe_decrypt(k).await.map_err(|e| {
324 let msg = format!("idx={} maybe_decrypt failed (raw_len={}): {}", idx, raw_len, e);
325 veilid_log!(self warn "TableDB::get_keys({}) col={} {}", self.unlocked_inner.table, col, msg);
326 VeilidAPIError::generic(msg)
327 })?;
328 let decompressed = decompress_size_prepended(&decrypted, max_value_size).map_err(|e| {
329 let preview: Vec<u8> = decrypted.as_ref()[..decrypted.len().min(16)].to_vec();
330 let msg = format!("idx={} decompress failed (raw_len={} decrypted_len={} first_bytes={:02x?}): {}", idx, raw_len, decrypted.len(), preview, e);
331 veilid_log!(self warn "TableDB::get_keys({}) col={} {}", self.unlocked_inner.table, col, msg);
332 std::io::Error::other(msg)
333 })?;
334 *k = decompressed;
335 }
336 Ok(out)
337 }
338
339 #[cfg_attr(
345 feature = "instrument",
346 instrument(level = "trace", target = "tstore", skip_all)
347 )]
348 pub async fn get_key_count(&self, col: u32) -> VeilidAPIResult<u64> {
349 if col >= self.opened_column_count {
350 apibail_generic!(
351 "Column exceeds opened column count {} >= {}",
352 col,
353 self.opened_column_count
354 );
355 }
356 let db = self.unlocked_inner.database.clone();
357 let key_count = db.num_keys(col).await.map_err(VeilidAPIError::from)?;
358 Ok(key_count)
359 }
360
361 #[cfg_attr(
365 feature = "instrument",
366 instrument(level = "trace", target = "tstore", skip_all)
367 )]
368 #[must_use]
369 pub fn transact(&self) -> TableDBTransaction {
370 let dbt = self.unlocked_inner.database.transaction();
371 TableDBTransaction::new(self.clone(), dbt)
372 }
373
374 #[cfg_attr(
380 feature = "instrument",
381 instrument(level = "trace", target = "tstore", skip_all)
382 )]
383 pub async fn store(&self, col: u32, key: &[u8], value: &[u8]) -> VeilidAPIResult<()> {
384 if col >= self.opened_column_count {
385 apibail_generic!(
386 "Column exceeds opened column count {} >= {}",
387 col,
388 self.opened_column_count
389 );
390 }
391 let db = self.unlocked_inner.database.clone();
392 let mut dbt = db.transaction();
393 dbt.put(
394 col,
395 self.maybe_encrypt(&compress_prepend_size(key), true).await,
396 self.maybe_encrypt(&compress_prepend_size(value), false)
397 .await,
398 );
399 db.write(dbt)
400 .await
401 .map_err(|e| VeilidAPIError::generic(format!("failed to store: {}", e)))
402 }
403
404 #[cfg_attr(
410 feature = "instrument",
411 instrument(level = "trace", target = "tstore", skip_all)
412 )]
413 pub async fn store_json<T>(&self, col: u32, key: &[u8], value: &T) -> VeilidAPIResult<()>
414 where
415 T: serde::Serialize,
416 {
417 let value = serde_json::to_vec(value).map_err(VeilidAPIError::internal)?;
418 self.store(col, key, &value).await
419 }
420
421 #[cfg_attr(
427 feature = "instrument",
428 instrument(level = "trace", target = "tstore", skip_all)
429 )]
430 pub async fn load(&self, col: u32, key: &[u8]) -> VeilidAPIResult<Option<Vec<u8>>> {
431 if col >= self.opened_column_count {
432 apibail_generic!(
433 "Column exceeds opened column count {} >= {}",
434 col,
435 self.opened_column_count
436 );
437 }
438 let db = self.unlocked_inner.database.clone();
439 let key = self.maybe_encrypt(&compress_prepend_size(key), true).await;
440 let max_value_size = self.config().table_store.max_value_size_mb as usize * 1024 * 1024;
441 match db.get(col, &key).await.map_err(VeilidAPIError::from)? {
442 Some(v) => Ok(Some(
443 decompress_size_prepended(
444 &self.maybe_decrypt(&v).await.map_err(VeilidAPIError::from)?,
445 max_value_size,
446 )
447 .map_err(|e| std::io::Error::other(e.to_string()))?,
448 )),
449 None => Ok(None),
450 }
451 }
452
453 #[cfg_attr(
459 feature = "instrument",
460 instrument(level = "trace", target = "tstore", skip_all)
461 )]
462 pub async fn load_json<T>(&self, col: u32, key: &[u8]) -> VeilidAPIResult<Option<T>>
463 where
464 T: for<'de> serde::Deserialize<'de>,
465 {
466 let out = match self.load(col, key).await? {
467 Some(v) => Some(serde_json::from_slice(&v).map_err(VeilidAPIError::internal)?),
468 None => None,
469 };
470 Ok(out)
471 }
472
473 #[cfg_attr(
479 feature = "instrument",
480 instrument(level = "trace", target = "tstore", skip_all)
481 )]
482 pub async fn delete(&self, col: u32, key: &[u8]) -> VeilidAPIResult<Option<Vec<u8>>> {
483 if col >= self.opened_column_count {
484 apibail_generic!(
485 "Column exceeds opened column count {} >= {}",
486 col,
487 self.opened_column_count
488 );
489 }
490 let key = self.maybe_encrypt(&compress_prepend_size(key), true).await;
491
492 let db = self.unlocked_inner.database.clone();
493
494 let max_value_size = self.config().table_store.max_value_size_mb as usize * 1024 * 1024;
495 match db.delete(col, &key).await.map_err(VeilidAPIError::from)? {
496 Some(v) => Ok(Some(
497 decompress_size_prepended(
498 &self.maybe_decrypt(&v).await.map_err(VeilidAPIError::from)?,
499 max_value_size,
500 )
501 .map_err(|e| std::io::Error::other(e.to_string()))?,
502 )),
503 None => Ok(None),
504 }
505 }
506
507 #[cfg_attr(
513 feature = "instrument",
514 instrument(level = "trace", target = "tstore", skip_all)
515 )]
516 pub async fn delete_json<T>(&self, col: u32, key: &[u8]) -> VeilidAPIResult<Option<T>>
517 where
518 T: for<'de> serde::Deserialize<'de>,
519 {
520 let old_value = match self.delete(col, key).await? {
521 Some(v) => Some(serde_json::from_slice(&v).map_err(VeilidAPIError::internal)?),
522 None => None,
523 };
524 Ok(old_value)
525 }
526}
527
528struct TableDBTransactionInner {
531 registry: VeilidComponentRegistry,
532 dbt: Option<DBTransaction>,
533}
534
535impl fmt::Debug for TableDBTransactionInner {
536 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
537 write!(
538 f,
539 "TableDBTransactionInner({})",
540 match &self.dbt {
541 Some(dbt) => format!("len={}", dbt.ops.len()),
542 None => "".to_owned(),
543 }
544 )
545 }
546}
547
548impl Drop for TableDBTransactionInner {
549 fn drop(&mut self) {
550 if self.dbt.is_some() {
551 let registry = &self.registry;
552 veilid_log!(registry error "Dropped transaction without commit or rollback");
553 }
554 }
555}
556
557#[derive(Debug, Clone)]
560pub struct TableDBTransaction {
561 db: TableDB,
562 inner: Arc<Mutex<TableDBTransactionInner>>,
563}
564
565impl VeilidComponentRegistryAccessor for TableDBTransaction {
566 fn registry(&self) -> VeilidComponentRegistry {
567 self.db.registry()
568 }
569}
570
571impl TableDBTransaction {
572 fn new(db: TableDB, dbt: DBTransaction) -> Self {
573 let registry = db.registry();
574 Self {
575 db,
576 inner: Arc::new(Mutex::new(TableDBTransactionInner {
577 registry,
578 dbt: Some(dbt),
579 })),
580 }
581 }
582
583 #[cfg_attr(
589 feature = "instrument",
590 instrument(level = "trace", target = "tstore", skip_all)
591 )]
592 pub async fn commit(self) -> VeilidAPIResult<()> {
593 let dbt = {
594 let mut inner = self.inner.lock();
595 inner
596 .dbt
597 .take()
598 .ok_or_else(|| VeilidAPIError::generic("transaction already completed"))?
599 };
600
601 if dbt.ops.is_empty() {
602 return Ok(());
604 }
605
606 let db = self.db.unlocked_inner.database.clone();
607 let _commit_lock = self
608 .db
609 .unlocked_inner
610 .commit_lock
611 .lock()
612 .measure_debug(
613 TimestampDuration::new_ms(200),
614 veilid_log_dbg!(
615 self,
616 "TableDBTransaction({})::commit lock",
617 self.db.table_name()
618 ),
619 )
620 .await;
621 db.write(dbt).await.map_err(|e| {
622 veilid_log!(self error "commit failed, transaction lost: {:?}", e);
623 VeilidAPIError::generic(format!("commit failed, transaction lost: {}", e))
624 })
625 }
626
627 #[cfg_attr(
631 feature = "instrument",
632 instrument(level = "trace", target = "tstore", skip_all)
633 )]
634 pub fn rollback(self) {
635 let mut inner = self.inner.lock();
636 inner.dbt = None;
637 }
638
639 #[cfg_attr(
645 feature = "instrument",
646 instrument(level = "trace", target = "tstore", skip_all)
647 )]
648 pub async fn store(&self, col: u32, key: &[u8], value: &[u8]) -> VeilidAPIResult<()> {
649 if col >= self.db.opened_column_count {
650 apibail_generic!(
651 "Column exceeds opened column count {} >= {}",
652 col,
653 self.db.opened_column_count
654 );
655 }
656
657 let key = self
658 .db
659 .maybe_encrypt(&compress_prepend_size(key), true)
660 .await;
661 let value = self
662 .db
663 .maybe_encrypt(&compress_prepend_size(value), false)
664 .await;
665 let mut inner = self.inner.lock();
666 inner
667 .dbt
668 .as_mut()
669 .ok_or_else(|| VeilidAPIError::generic("store failed, transaction already completed"))?
670 .put_owned(col, key.to_vec(), value.to_vec());
671 Ok(())
672 }
673
674 #[cfg_attr(
680 feature = "instrument",
681 instrument(level = "trace", target = "tstore", skip_all)
682 )]
683 pub async fn store_json<T>(&self, col: u32, key: &[u8], value: &T) -> VeilidAPIResult<()>
684 where
685 T: serde::Serialize,
686 {
687 let value = serde_json::to_vec(value).map_err(VeilidAPIError::internal)?;
688 self.store(col, key, &value).await
689 }
690
691 #[cfg_attr(
697 feature = "instrument",
698 instrument(level = "trace", target = "tstore", skip_all)
699 )]
700 pub async fn delete(&self, col: u32, key: &[u8]) -> VeilidAPIResult<()> {
701 if col >= self.db.opened_column_count {
702 apibail_generic!(
703 "Column exceeds opened column count {} >= {}",
704 col,
705 self.db.opened_column_count
706 );
707 }
708
709 let key = self
710 .db
711 .maybe_encrypt(&compress_prepend_size(key), true)
712 .await;
713 let mut inner = self.inner.lock();
714 inner
715 .dbt
716 .as_mut()
717 .ok_or_else(|| VeilidAPIError::generic("delete failed, transaction already completed"))?
718 .delete_owned(col, key.to_vec());
719 Ok(())
720 }
721}