1use std::collections::HashMap;
23use std::sync::atomic::{AtomicU64, Ordering};
24use std::sync::{Arc, Mutex};
25use std::time::{SystemTime, UNIX_EPOCH};
26
27use crate::runtime::RedDBRuntime;
28use crate::storage::query::engine::cancel::CancelToken;
29use reddb_types::types::Value;
30
31const RED_CONFIG_COLLECTION: &str = "red_config";
32
33pub const PAGE_SIZE: usize = 16 * 1024;
36
37pub trait Clock: Send + Sync {
40 fn now_ms(&self) -> u64;
41}
42
43#[derive(Debug, Default)]
44pub struct SystemClock;
45
46impl Clock for SystemClock {
47 fn now_ms(&self) -> u64 {
48 SystemTime::now()
49 .duration_since(UNIX_EPOCH)
50 .map(|d| d.as_millis() as u64)
51 .unwrap_or(0)
52 }
53}
54
55#[derive(Debug)]
56pub struct FakeClock {
57 now_ms: AtomicU64,
58}
59
60impl FakeClock {
61 pub fn new(now_ms: u64) -> Self {
62 Self {
63 now_ms: AtomicU64::new(now_ms),
64 }
65 }
66
67 pub fn advance(&self, ms: u64) {
68 self.now_ms.fetch_add(ms, Ordering::SeqCst);
69 }
70}
71
72impl Clock for FakeClock {
73 fn now_ms(&self) -> u64 {
74 self.now_ms.load(Ordering::SeqCst)
75 }
76}
77
78#[derive(Debug, Clone, Copy)]
83pub struct StreamConfig {
84 pub snapshot_ttl_ms: u64,
85 pub chunk_default_pages: usize,
86 pub chunk_min_pages: usize,
87 pub chunk_max_pages: usize,
88 pub chunk_max_rows: usize,
89 pub chunk_max_latency_ms: u64,
90 pub max_global_streams: usize,
94 pub max_per_principal_streams: usize,
99 pub default_verify: crate::runtime::integrity_tombstone::VerifyMode,
103}
104
105impl Default for StreamConfig {
106 fn default() -> Self {
107 Self::DEFAULT
108 }
109}
110
111impl StreamConfig {
112 pub const DEFAULT: StreamConfig = StreamConfig {
115 snapshot_ttl_ms: 60_000,
116 chunk_default_pages: 4,
117 chunk_min_pages: 1,
118 chunk_max_pages: 64,
119 chunk_max_rows: 1000,
120 chunk_max_latency_ms: 50,
121 max_global_streams: 256,
122 max_per_principal_streams: 32,
123 default_verify: crate::runtime::integrity_tombstone::VerifyMode::None,
124 };
125
126 pub fn load(runtime: &RedDBRuntime) -> Self {
131 let db = runtime.db();
132 let read_u64 = |key: &str| -> Option<u64> {
133 match db.get_kv(RED_CONFIG_COLLECTION, key) {
134 Some((Value::Integer(v), _)) if v >= 0 => Some(v as u64),
135 Some((Value::UnsignedInteger(v), _)) => Some(v),
136 Some((Value::Text(text), _)) => text.parse().ok(),
137 _ => None,
138 }
139 };
140
141 let mut cfg = Self::DEFAULT;
142 if let Some(v) = read_u64("stream.snapshot.ttl_ms") {
143 cfg.snapshot_ttl_ms = v;
144 }
145 if let Some(v) = read_u64("stream.chunk.default_pages") {
146 cfg.chunk_default_pages = v as usize;
147 }
148 if let Some(v) = read_u64("stream.chunk.min_pages") {
149 cfg.chunk_min_pages = v as usize;
150 }
151 if let Some(v) = read_u64("stream.chunk.max_pages") {
152 cfg.chunk_max_pages = v as usize;
153 }
154 if let Some(v) = read_u64("stream.chunk.max_rows") {
155 cfg.chunk_max_rows = v as usize;
156 }
157 if let Some(v) = read_u64("stream.chunk.max_latency_ms") {
158 cfg.chunk_max_latency_ms = v;
159 }
160 if let Some(v) = read_u64("stream.max_global") {
161 cfg.max_global_streams = v as usize;
162 }
163 if let Some(v) = read_u64("stream.max_per_principal") {
164 cfg.max_per_principal_streams = v as usize;
165 }
166 if let Some((Value::Text(text), _)) =
169 db.get_kv(RED_CONFIG_COLLECTION, "stream.integrity.default_verify")
170 {
171 cfg.default_verify = crate::runtime::integrity_tombstone::VerifyMode::parse(&text);
172 }
173 cfg.normalize();
174 cfg
175 }
176
177 fn normalize(&mut self) {
181 if self.chunk_min_pages == 0 {
182 self.chunk_min_pages = 1;
183 }
184 if self.chunk_max_pages < self.chunk_min_pages {
185 self.chunk_max_pages = self.chunk_min_pages;
186 }
187 if self.chunk_default_pages < self.chunk_min_pages {
188 self.chunk_default_pages = self.chunk_min_pages;
189 }
190 if self.chunk_default_pages > self.chunk_max_pages {
191 self.chunk_default_pages = self.chunk_max_pages;
192 }
193 if self.chunk_max_rows == 0 {
194 self.chunk_max_rows = 1;
195 }
196 if self.max_global_streams == 0 {
197 self.max_global_streams = 1;
198 }
199 if self.max_per_principal_streams == 0 {
200 self.max_per_principal_streams = 1;
201 }
202 }
203
204 pub fn production_buffer_bytes(&self) -> usize {
207 self.chunk_default_pages.saturating_mul(PAGE_SIZE)
208 }
209}
210
211static NEXT_LEASE_ID: AtomicU64 = AtomicU64::new(1);
219
220pub const LEASE_HANDLE_BYTES: usize = 16;
226
227pub fn generate_lease_handle() -> String {
231 let mut bytes = [0u8; LEASE_HANDLE_BYTES];
232 if crate::crypto::os_random::fill_bytes(&mut bytes).is_err() {
233 let lo = NEXT_LEASE_ID.fetch_add(0, Ordering::SeqCst).to_le_bytes();
239 let now = crate::utils::now_unix_nanos().to_le_bytes();
240 bytes[..8].copy_from_slice(&lo);
241 bytes[8..].copy_from_slice(&now);
242 }
243 crate::utils::to_hex(&bytes)
244}
245
246pub fn parse_jwt_exp_ms(token: &str) -> Option<u64> {
256 let parts: Vec<&str> = token.split('.').collect();
257 if parts.len() != 3 {
258 return None;
259 }
260 let payload = base64url_decode_padded(parts[1])?;
261 let v: crate::json::Value = crate::json::from_slice(&payload).ok()?;
262 let exp_secs = v.get("exp").and_then(|n| n.as_f64())? as i64;
263 if exp_secs <= 0 {
264 return None;
265 }
266 Some((exp_secs as u64).saturating_mul(1000))
267}
268
269fn base64url_decode_padded(input: &str) -> Option<Vec<u8>> {
270 let mut s = String::with_capacity(input.len() + 4);
271 for ch in input.chars() {
272 match ch {
273 '-' => s.push('+'),
274 '_' => s.push('/'),
275 _ => s.push(ch),
276 }
277 }
278 while !s.len().is_multiple_of(4) {
279 s.push('=');
280 }
281 reddb_wire::redwire::handshake::base64_std_decode(&s)
282}
283
284#[derive(Debug, Clone, Copy, PartialEq, Eq)]
288pub enum CloseReason {
289 Ok,
290 Cancelled,
291 Error,
292 SnapshotExpired,
293 CapacityRefused,
294 IntegrityFailed,
295}
296
297impl CloseReason {
298 pub fn as_str(&self) -> &'static str {
299 match self {
300 CloseReason::Ok => "ok",
301 CloseReason::Cancelled => "cancelled",
302 CloseReason::Error => "error",
303 CloseReason::SnapshotExpired => "snapshot_expired",
304 CloseReason::CapacityRefused => "capacity_refused",
305 CloseReason::IntegrityFailed => "integrity_failed",
306 }
307 }
308}
309
310#[derive(Debug)]
311pub struct StreamLease {
312 pub id: u64,
313 pub lease_handle: String,
317 pub snapshot_lsn: u64,
318 pub opened_at_ms: u64,
319 pub config: StreamConfig,
320}
321
322impl StreamLease {
323 pub fn snapshot_expired(&self, now_ms: u64) -> bool {
330 now_ms.saturating_sub(self.opened_at_ms) >= self.config.snapshot_ttl_ms
331 }
332}
333
334#[derive(Debug, PartialEq, Eq)]
335pub enum OpenStreamError {
336 TransactionActive,
340}
341
342impl OpenStreamError {
343 pub fn code(&self) -> &'static str {
344 match self {
345 OpenStreamError::TransactionActive => "stream_in_transaction_unsupported",
346 }
347 }
348
349 pub fn message(&self) -> &'static str {
350 match self {
351 OpenStreamError::TransactionActive => {
352 "cannot open output stream while a transaction is active on this session"
353 }
354 }
355 }
356}
357
358pub fn open_stream(
363 config: StreamConfig,
364 snapshot_lsn: u64,
365 in_transaction: bool,
366 clock: &dyn Clock,
367) -> Result<StreamLease, OpenStreamError> {
368 if in_transaction {
369 return Err(OpenStreamError::TransactionActive);
370 }
371 Ok(StreamLease {
372 id: NEXT_LEASE_ID.fetch_add(1, Ordering::SeqCst),
373 lease_handle: generate_lease_handle(),
374 snapshot_lsn,
375 opened_at_ms: clock.now_ms(),
376 config,
377 })
378}
379
380pub struct ChunkProducer<'a> {
388 buf: Vec<u8>,
389 rows: usize,
390 window_started_ms: u64,
391 cap_bytes: usize,
392 cap_rows: usize,
393 cap_latency_ms: u64,
394 clock: &'a dyn Clock,
395 total_flushes: u64,
396 total_bytes: u64,
397 total_rows: u64,
398 last_flush_reason: Option<FlushReason>,
399}
400
401#[derive(Debug, Clone, Copy, PartialEq, Eq)]
402pub enum FlushReason {
403 Byte,
404 Row,
405 Latency,
406 Terminal,
407}
408
409impl<'a> ChunkProducer<'a> {
410 pub fn new(config: &StreamConfig, clock: &'a dyn Clock) -> Self {
411 let cap_bytes = config.production_buffer_bytes();
412 Self {
413 buf: Vec::with_capacity(cap_bytes),
414 rows: 0,
415 window_started_ms: clock.now_ms(),
416 cap_bytes,
417 cap_rows: config.chunk_max_rows,
418 cap_latency_ms: config.chunk_max_latency_ms,
419 clock,
420 total_flushes: 0,
421 total_bytes: 0,
422 total_rows: 0,
423 last_flush_reason: None,
424 }
425 }
426
427 pub fn push_line<F>(&mut self, line: &[u8], flush: &mut F) -> std::io::Result<bool>
430 where
431 F: FnMut(&[u8]) -> std::io::Result<()>,
432 {
433 crate::runtime::query_budget::charge_output(line.len()).map_err(std::io::Error::other)?;
434 self.buf.extend_from_slice(line);
435 self.rows += 1;
436 self.total_rows += 1;
437
438 if self.buf.len() >= self.cap_bytes {
439 self.flush(flush, FlushReason::Byte)?;
440 return Ok(true);
441 }
442 if self.rows >= self.cap_rows {
443 self.flush(flush, FlushReason::Row)?;
444 return Ok(true);
445 }
446 let elapsed = self.clock.now_ms().saturating_sub(self.window_started_ms);
447 if elapsed >= self.cap_latency_ms {
448 self.flush(flush, FlushReason::Latency)?;
449 return Ok(true);
450 }
451 Ok(false)
452 }
453
454 pub fn drive_lines<S, R, Enc, F>(
472 &mut self,
473 source: S,
474 mut encode: Enc,
475 flush: &mut F,
476 ) -> std::io::Result<u64>
477 where
478 S: IntoIterator<Item = R>,
479 Enc: FnMut(&R) -> Vec<u8>,
480 F: FnMut(&[u8]) -> std::io::Result<()>,
481 {
482 let mut count = 0u64;
483 for record in source {
484 let line = encode(&record);
485 self.push_line(&line, flush)?;
486 count += 1;
487 }
488 Ok(count)
489 }
490
491 pub fn finish<F>(&mut self, flush: &mut F) -> std::io::Result<()>
495 where
496 F: FnMut(&[u8]) -> std::io::Result<()>,
497 {
498 if !self.buf.is_empty() {
499 self.flush(flush, FlushReason::Terminal)?;
500 }
501 Ok(())
502 }
503
504 fn flush<F>(&mut self, flush: &mut F, reason: FlushReason) -> std::io::Result<()>
505 where
506 F: FnMut(&[u8]) -> std::io::Result<()>,
507 {
508 flush(&self.buf)?;
509 self.total_bytes += self.buf.len() as u64;
510 self.total_flushes += 1;
511 self.last_flush_reason = Some(reason);
512 self.buf.clear();
513 self.rows = 0;
514 self.window_started_ms = self.clock.now_ms();
515 Ok(())
516 }
517
518 pub fn total_flushes(&self) -> u64 {
519 self.total_flushes
520 }
521 pub fn total_bytes(&self) -> u64 {
522 self.total_bytes
523 }
524 pub fn total_rows(&self) -> u64 {
525 self.total_rows
526 }
527 pub fn last_flush_reason(&self) -> Option<FlushReason> {
528 self.last_flush_reason
529 }
530}
531
532pub fn write_chunked_response_header<W: std::io::Write>(
536 writer: &mut W,
537 status: u16,
538 content_type: &str,
539) -> std::io::Result<()> {
540 let mut header = format!(
548 "HTTP/1.1 {} {}\r\nContent-Type: {}\r\nTransfer-Encoding: chunked\r\nCache-Control: no-cache\r\nConnection: close\r\n",
549 status,
550 crate::server::transport::status_text(status),
551 content_type,
552 );
553 for (name, value) in crate::server::transport::cors_header_pairs() {
557 header.push_str(name);
558 header.push_str(": ");
559 header.push_str(&value);
560 header.push_str("\r\n");
561 }
562 header.push_str("\r\n");
563 writer.write_all(header.as_bytes())?;
564 writer.flush()
565}
566
567pub fn write_chunk<W: std::io::Write>(writer: &mut W, bytes: &[u8]) -> std::io::Result<()> {
571 if bytes.is_empty() {
572 return Ok(());
573 }
574 let size = format!("{:x}\r\n", bytes.len());
575 writer.write_all(size.as_bytes())?;
576 writer.write_all(bytes)?;
577 writer.write_all(b"\r\n")?;
578 writer.flush()
579}
580
581pub fn write_chunked_terminator<W: std::io::Write>(writer: &mut W) -> std::io::Result<()> {
583 writer.write_all(b"0\r\n\r\n")?;
584 writer.flush()
585}
586
587#[derive(Debug, Default)]
595pub struct StreamCapacityRegistry {
596 inner: Mutex<CapacityInner>,
597}
598
599#[derive(Debug, Default)]
600struct CapacityInner {
601 global_count: usize,
602 per_principal: HashMap<String, usize>,
603}
604
605#[derive(Debug, Clone, PartialEq, Eq)]
610pub enum AcquireError {
611 GlobalExhausted { limit: usize, current: usize },
613 PrincipalExhausted {
617 principal: String,
618 limit: usize,
619 current: usize,
620 },
621}
622
623impl AcquireError {
624 pub fn code(&self) -> &'static str {
625 match self {
626 AcquireError::GlobalExhausted { .. } => "server_stream_capacity_exhausted",
627 AcquireError::PrincipalExhausted { .. } => "principal_stream_quota_exhausted",
628 }
629 }
630
631 pub fn message(&self) -> String {
632 match self {
633 AcquireError::GlobalExhausted { limit, current } => {
634 format!("server stream capacity exhausted (limit {limit}, current {current})")
635 }
636 AcquireError::PrincipalExhausted {
637 principal,
638 limit,
639 current,
640 } => format!(
641 "principal {principal} stream quota exhausted (limit {limit}, current {current})"
642 ),
643 }
644 }
645}
646
647impl StreamCapacityRegistry {
648 pub fn new() -> Arc<Self> {
649 Arc::new(Self::default())
650 }
651
652 pub fn try_acquire(
656 self: &Arc<Self>,
657 principal: &str,
658 max_global: usize,
659 max_per_principal: usize,
660 ) -> Result<StreamCapacityGuard, AcquireError> {
661 let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
662 if inner.global_count >= max_global {
663 return Err(AcquireError::GlobalExhausted {
664 limit: max_global,
665 current: inner.global_count,
666 });
667 }
668 let current = inner.per_principal.get(principal).copied().unwrap_or(0);
669 if current >= max_per_principal {
670 return Err(AcquireError::PrincipalExhausted {
671 principal: principal.to_string(),
672 limit: max_per_principal,
673 current,
674 });
675 }
676 inner.global_count += 1;
677 inner
678 .per_principal
679 .insert(principal.to_string(), current + 1);
680 Ok(StreamCapacityGuard {
681 registry: Arc::clone(self),
682 principal: principal.to_string(),
683 released: false,
684 })
685 }
686
687 fn release(&self, principal: &str) {
688 let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
689 if inner.global_count > 0 {
690 inner.global_count -= 1;
691 }
692 if let Some(count) = inner.per_principal.get_mut(principal) {
693 if *count > 0 {
694 *count -= 1;
695 }
696 if *count == 0 {
697 inner.per_principal.remove(principal);
698 }
699 }
700 }
701
702 pub fn snapshot(&self) -> (usize, HashMap<String, usize>) {
704 let inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
705 (inner.global_count, inner.per_principal.clone())
706 }
707}
708
709#[must_use = "dropping the guard immediately releases the stream slot"]
714#[derive(Debug)]
715pub struct StreamCapacityGuard {
716 registry: Arc<StreamCapacityRegistry>,
717 principal: String,
718 released: bool,
719}
720
721impl StreamCapacityGuard {
722 pub fn principal(&self) -> &str {
723 &self.principal
724 }
725}
726
727impl Drop for StreamCapacityGuard {
728 fn drop(&mut self) {
729 if !self.released {
730 self.registry.release(&self.principal);
731 self.released = true;
732 }
733 }
734}
735
736pub fn assess_resumability(query: &str) -> bool {
745 let upper = query.to_uppercase();
746 let trimmed = upper.trim_start();
747 if !trimmed.starts_with("SELECT ") && !trimmed.starts_with("SELECT\n") {
748 return false;
749 }
750 const FORBIDDEN: &[&str] = &[
751 " GROUP BY ",
752 " HAVING ",
753 " DISTINCT ",
754 "DISTINCT ",
755 "COUNT(",
756 "SUM(",
757 "AVG(",
758 "MIN(",
759 "MAX(",
760 "ARRAY_AGG(",
761 "JSON_AGG(",
762 "OVER(",
763 " OVER (",
764 " JOIN ",
765 ];
766 for kw in FORBIDDEN {
767 if upper.contains(kw) {
768 return false;
769 }
770 }
771 if let Some(idx) = upper.find("ORDER BY") {
772 let tail = &upper[idx + "ORDER BY".len()..];
773 let mut clause = tail.to_string();
775 if let Some(lim) = clause.find(" LIMIT ") {
776 clause.truncate(lim);
777 }
778 if let Some(semi) = clause.find(';') {
779 clause.truncate(semi);
780 }
781 let clause = clause.trim();
782 if !matches!(clause, "RID" | "RID ASC") {
783 return false;
784 }
785 }
786 true
787}
788
789#[derive(Debug, Default)]
795pub struct LeaseRegistry {
796 inner: Mutex<HashMap<u64, LeaseEntry>>,
797}
798
799#[derive(Debug, Clone, Copy)]
800struct LeaseEntry {
801 opened_at_ms: u64,
802 ttl_ms: u64,
803}
804
805#[derive(Debug, Clone, Copy, PartialEq, Eq)]
806pub enum LeaseLookup {
807 Unknown,
809 Expired,
811 Live,
813}
814
815impl LeaseRegistry {
816 pub fn new() -> Arc<Self> {
817 Arc::new(Self::default())
818 }
819
820 pub fn record(&self, snapshot_lsn: u64, opened_at_ms: u64, ttl_ms: u64) {
825 let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
826 inner.insert(
827 snapshot_lsn,
828 LeaseEntry {
829 opened_at_ms,
830 ttl_ms,
831 },
832 );
833 }
834
835 pub fn lookup(&self, snapshot_lsn: u64, now_ms: u64) -> LeaseLookup {
838 let inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
839 match inner.get(&snapshot_lsn) {
840 None => LeaseLookup::Unknown,
841 Some(entry) => {
842 if now_ms.saturating_sub(entry.opened_at_ms) >= entry.ttl_ms {
843 LeaseLookup::Expired
844 } else {
845 LeaseLookup::Live
846 }
847 }
848 }
849 }
850
851 #[doc(hidden)]
853 pub fn len(&self) -> usize {
854 self.inner.lock().unwrap_or_else(|e| e.into_inner()).len()
855 }
856}
857
858pub const CURSOR_TOKEN_BYTES: usize = 24;
865
866static NEXT_CURSOR_SALT: AtomicU64 = AtomicU64::new(1);
867
868pub fn generate_cursor_token() -> String {
873 let mut bytes = [0u8; CURSOR_TOKEN_BYTES];
874 if crate::crypto::os_random::fill_bytes(&mut bytes).is_err() {
875 let salt = NEXT_CURSOR_SALT
876 .fetch_add(1, Ordering::SeqCst)
877 .to_le_bytes();
878 let now = crate::utils::now_unix_nanos().to_le_bytes();
879 bytes[..8].copy_from_slice(&salt);
880 bytes[8..16].copy_from_slice(&now);
881 }
882 crate::utils::to_hex(&bytes)
883}
884
885#[derive(Debug, Clone)]
889struct CursorEntry {
890 snapshot_lsn: u64,
891 tenant: String,
892 principal: String,
893 query: String,
894 entity_types: Option<Vec<String>>,
895 capabilities: Option<Vec<String>>,
896 opened_at_ms: u64,
897 ttl_ms: u64,
898 cancel_token: CancelToken,
903 cancelled: bool,
907}
908
909#[derive(Debug, Clone, PartialEq, Eq)]
913pub struct CursorResume {
914 pub snapshot_lsn: u64,
915 pub query: String,
916 pub entity_types: Option<Vec<String>>,
917 pub capabilities: Option<Vec<String>>,
918 pub expires_at_ms: u64,
919}
920
921#[derive(Debug, Clone, Copy, PartialEq, Eq)]
926pub enum CursorReject {
927 NotFound,
930 Expired,
932 Cancelled,
937}
938
939#[derive(Debug, Default)]
950pub struct CursorRegistry {
951 inner: Mutex<HashMap<String, CursorEntry>>,
952}
953
954impl CursorRegistry {
955 pub fn new() -> Arc<Self> {
956 Arc::new(Self::default())
957 }
958
959 #[allow(clippy::too_many_arguments)]
963 pub fn register(
964 &self,
965 snapshot_lsn: u64,
966 tenant: &str,
967 principal: &str,
968 query: &str,
969 entity_types: Option<Vec<String>>,
970 capabilities: Option<Vec<String>>,
971 opened_at_ms: u64,
972 ttl_ms: u64,
973 ) -> String {
974 let token = generate_cursor_token();
975 let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
976 inner.insert(
977 token.clone(),
978 CursorEntry {
979 snapshot_lsn,
980 tenant: tenant.to_string(),
981 principal: principal.to_string(),
982 query: query.to_string(),
983 entity_types,
984 capabilities,
985 opened_at_ms,
986 ttl_ms,
987 cancel_token: CancelToken::new(),
988 cancelled: false,
989 },
990 );
991 token
992 }
993
994 pub fn cancel_token_for(&self, token: &str) -> Option<CancelToken> {
999 let inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
1000 inner.get(token).map(|e| e.cancel_token.clone())
1001 }
1002
1003 pub fn cancel(
1014 &self,
1015 token: &str,
1016 tenant: &str,
1017 principal: &str,
1018 ) -> Result<CancelToken, CursorReject> {
1019 let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
1020 let entry = inner.get_mut(token).ok_or(CursorReject::NotFound)?;
1021 if entry.tenant != tenant || entry.principal != principal {
1022 return Err(CursorReject::NotFound);
1023 }
1024 entry.cancelled = true;
1025 entry.cancel_token.cancel();
1026 Ok(entry.cancel_token.clone())
1027 }
1028
1029 pub fn resolve(
1037 &self,
1038 token: &str,
1039 tenant: &str,
1040 principal: &str,
1041 now_ms: u64,
1042 ) -> Result<CursorResume, CursorReject> {
1043 let inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
1044 let entry = inner.get(token).ok_or(CursorReject::NotFound)?;
1045 if entry.tenant != tenant || entry.principal != principal {
1049 return Err(CursorReject::NotFound);
1050 }
1051 if entry.cancelled {
1055 return Err(CursorReject::Cancelled);
1056 }
1057 if now_ms.saturating_sub(entry.opened_at_ms) >= entry.ttl_ms {
1058 return Err(CursorReject::Expired);
1059 }
1060 Ok(CursorResume {
1061 snapshot_lsn: entry.snapshot_lsn,
1062 query: entry.query.clone(),
1063 entity_types: entry.entity_types.clone(),
1064 capabilities: entry.capabilities.clone(),
1065 expires_at_ms: entry.opened_at_ms.saturating_add(entry.ttl_ms),
1066 })
1067 }
1068
1069 #[doc(hidden)]
1071 pub fn len(&self) -> usize {
1072 self.inner.lock().unwrap_or_else(|e| e.into_inner()).len()
1073 }
1074
1075 #[doc(hidden)]
1076 pub fn is_empty(&self) -> bool {
1077 self.len() == 0
1078 }
1079}
1080
1081#[derive(Debug, Default)]
1087pub struct PrefixHasher {
1088 inner: Option<sha2::Sha256>,
1089 rows: u64,
1090}
1091
1092impl PrefixHasher {
1093 pub fn new() -> Self {
1094 use sha2::Digest;
1095 Self {
1096 inner: Some(sha2::Sha256::new()),
1097 rows: 0,
1098 }
1099 }
1100
1101 pub fn update(&mut self, line: &[u8]) {
1102 use sha2::Digest;
1103 if let Some(h) = self.inner.as_mut() {
1104 h.update(line);
1105 }
1106 self.rows += 1;
1107 }
1108
1109 pub fn rows(&self) -> u64 {
1110 self.rows
1111 }
1112
1113 pub fn finalize_hex(mut self) -> String {
1116 use sha2::Digest;
1117 let hasher = self
1118 .inner
1119 .take()
1120 .expect("PrefixHasher::finalize_hex called twice");
1121 let digest = hasher.finalize();
1122 let mut out = String::with_capacity(64);
1123 for b in digest.iter() {
1124 out.push_str(&format!("{b:02x}"));
1125 }
1126 out
1127 }
1128}
1129
1130pub fn audit_stream_opened(
1136 runtime: &RedDBRuntime,
1137 lease_handle: &str,
1138 principal: &str,
1139 snapshot_lsn: u64,
1140 query_hash: &str,
1141) {
1142 use crate::json::{Map, Value as JsonValue};
1143 let mut detail = Map::new();
1144 detail.insert(
1145 "lease_handle".to_string(),
1146 JsonValue::String(lease_handle.to_string()),
1147 );
1148 detail.insert(
1149 "snapshot_lsn".to_string(),
1150 JsonValue::Number(snapshot_lsn as f64),
1151 );
1152 detail.insert(
1153 "query_hash".to_string(),
1154 JsonValue::String(query_hash.to_string()),
1155 );
1156 let event = crate::runtime::audit_log::AuditEvent::builder("stream.opened")
1157 .principal(principal)
1158 .resource(lease_handle.to_string())
1159 .outcome(crate::runtime::audit_log::Outcome::Success)
1160 .detail(JsonValue::Object(detail))
1161 .build();
1162 runtime.audit_log().record_event(event);
1163}
1164
1165pub fn audit_stream_closed(
1166 runtime: &RedDBRuntime,
1167 lease_handle: &str,
1168 principal: &str,
1169 reason: CloseReason,
1170 row_count: u64,
1171 bytes_written: u64,
1172) {
1173 use crate::json::{Map, Value as JsonValue};
1174 let mut stats = Map::new();
1175 stats.insert("row_count".to_string(), JsonValue::Number(row_count as f64));
1176 stats.insert(
1177 "bytes_written".to_string(),
1178 JsonValue::Number(bytes_written as f64),
1179 );
1180 let mut detail = Map::new();
1181 detail.insert(
1182 "lease_handle".to_string(),
1183 JsonValue::String(lease_handle.to_string()),
1184 );
1185 detail.insert(
1186 "reason".to_string(),
1187 JsonValue::String(reason.as_str().to_string()),
1188 );
1189 detail.insert("stats".to_string(), JsonValue::Object(stats));
1190 let outcome = match reason {
1191 CloseReason::Ok => crate::runtime::audit_log::Outcome::Success,
1192 CloseReason::CapacityRefused => crate::runtime::audit_log::Outcome::Denied,
1193 _ => crate::runtime::audit_log::Outcome::Error,
1194 };
1195 let event = crate::runtime::audit_log::AuditEvent::builder("stream.closed")
1196 .principal(principal)
1197 .resource(lease_handle.to_string())
1198 .outcome(outcome)
1199 .detail(JsonValue::Object(detail))
1200 .build();
1201 runtime.audit_log().record_event(event);
1202}
1203
1204pub fn audit_token_expired_during_lease(
1205 runtime: &RedDBRuntime,
1206 lease_handle: &str,
1207 principal: &str,
1208 token_expiry_ms: u64,
1209) {
1210 use crate::json::{Map, Value as JsonValue};
1211 let mut detail = Map::new();
1212 detail.insert(
1213 "lease_handle".to_string(),
1214 JsonValue::String(lease_handle.to_string()),
1215 );
1216 detail.insert(
1217 "token_expiry".to_string(),
1218 JsonValue::Number(token_expiry_ms as f64),
1219 );
1220 detail.insert("lease_continued".to_string(), JsonValue::Bool(true));
1221 let event = crate::runtime::audit_log::AuditEvent::builder("stream.token_expired_during_lease")
1222 .principal(principal)
1223 .resource(lease_handle.to_string())
1224 .outcome(crate::runtime::audit_log::Outcome::Success)
1225 .detail(JsonValue::Object(detail))
1226 .build();
1227 runtime.audit_log().record_event(event);
1228}
1229
1230pub fn audit_stream_capacity_refused(
1236 runtime: &RedDBRuntime,
1237 principal: &str,
1238 code: &str,
1239 limit: usize,
1240 current: usize,
1241) {
1242 use crate::json::{Map, Value as JsonValue};
1243 let mut detail = Map::new();
1244 detail.insert(
1245 "reason".to_string(),
1246 JsonValue::String(CloseReason::CapacityRefused.as_str().to_string()),
1247 );
1248 detail.insert("code".to_string(), JsonValue::String(code.to_string()));
1249 detail.insert("limit".to_string(), JsonValue::Number(limit as f64));
1250 detail.insert("current".to_string(), JsonValue::Number(current as f64));
1251 let event = crate::runtime::audit_log::AuditEvent::builder("stream.closed")
1252 .principal(principal)
1253 .outcome(crate::runtime::audit_log::Outcome::Denied)
1254 .detail(JsonValue::Object(detail))
1255 .build();
1256 runtime.audit_log().record_event(event);
1257}
1258
1259pub fn system_clock() -> Arc<dyn Clock> {
1262 static INSTANCE: std::sync::OnceLock<Arc<dyn Clock>> = std::sync::OnceLock::new();
1263 Arc::clone(INSTANCE.get_or_init(|| Arc::new(SystemClock)))
1264}
1265
1266#[cfg(test)]
1267mod tests {
1268 use super::*;
1269
1270 #[test]
1271 fn open_stream_refuses_when_session_has_active_transaction() {
1272 let clock = FakeClock::new(0);
1273 let err = open_stream(StreamConfig::DEFAULT, 42, true, &clock).unwrap_err();
1274 assert_eq!(err, OpenStreamError::TransactionActive);
1275 assert_eq!(err.code(), "stream_in_transaction_unsupported");
1276 }
1277
1278 #[test]
1279 fn open_stream_succeeds_when_session_is_autocommit() {
1280 let clock = FakeClock::new(1_700_000_000_000);
1281 let lease = open_stream(StreamConfig::DEFAULT, 123, false, &clock).unwrap();
1282 assert_eq!(lease.snapshot_lsn, 123);
1283 assert_eq!(lease.opened_at_ms, 1_700_000_000_000);
1284 assert!(lease.id >= 1);
1285 }
1286
1287 #[test]
1288 fn lease_ids_are_unique_and_monotonic() {
1289 let clock = FakeClock::new(0);
1290 let a = open_stream(StreamConfig::DEFAULT, 1, false, &clock).unwrap();
1291 let b = open_stream(StreamConfig::DEFAULT, 1, false, &clock).unwrap();
1292 assert!(b.id > a.id);
1293 }
1294
1295 #[test]
1296 fn snapshot_expired_uses_injected_clock_and_ttl() {
1297 let clock = FakeClock::new(0);
1300 let mut config = StreamConfig::DEFAULT;
1301 config.snapshot_ttl_ms = 5_000;
1302 let lease = open_stream(config, 0, false, &clock).unwrap();
1303
1304 assert!(!lease.snapshot_expired(clock.now_ms()));
1305 clock.advance(4_999);
1306 assert!(!lease.snapshot_expired(clock.now_ms()));
1307 clock.advance(1);
1308 assert!(lease.snapshot_expired(clock.now_ms()));
1309 }
1310
1311 #[test]
1312 fn stream_config_loads_defaults_when_kv_is_empty() {
1313 let cfg = StreamConfig::DEFAULT;
1315 assert_eq!(cfg.snapshot_ttl_ms, 60_000);
1316 assert_eq!(cfg.chunk_default_pages, 4);
1317 assert_eq!(cfg.chunk_min_pages, 1);
1318 assert_eq!(cfg.chunk_max_pages, 64);
1319 assert_eq!(cfg.chunk_max_rows, 1000);
1320 assert_eq!(cfg.chunk_max_latency_ms, 50);
1321 assert_eq!(cfg.production_buffer_bytes(), 64 * 1024);
1322 }
1323
1324 #[test]
1325 fn stream_config_normalize_clamps_inconsistent_inputs() {
1326 let mut cfg = StreamConfig {
1327 snapshot_ttl_ms: 1,
1328 chunk_default_pages: 100,
1329 chunk_min_pages: 0,
1330 chunk_max_pages: 8,
1331 chunk_max_rows: 0,
1332 chunk_max_latency_ms: 1,
1333 max_global_streams: 0,
1334 max_per_principal_streams: 0,
1335 default_verify: crate::runtime::integrity_tombstone::VerifyMode::None,
1336 };
1337 cfg.normalize();
1338 assert_eq!(cfg.chunk_min_pages, 1);
1339 assert_eq!(cfg.chunk_max_pages, 8);
1340 assert_eq!(cfg.chunk_default_pages, 8); assert!(cfg.chunk_max_rows >= 1);
1342 assert!(cfg.max_global_streams >= 1);
1343 assert!(cfg.max_per_principal_streams >= 1);
1344 }
1345
1346 struct CapturingSink {
1351 chunks: std::cell::RefCell<Vec<Vec<u8>>>,
1352 }
1353 impl CapturingSink {
1354 fn new() -> Self {
1355 Self {
1356 chunks: std::cell::RefCell::new(Vec::new()),
1357 }
1358 }
1359 fn len(&self) -> usize {
1360 self.chunks.borrow().len()
1361 }
1362 fn last_len(&self) -> Option<usize> {
1363 self.chunks.borrow().last().map(|c| c.len())
1364 }
1365 }
1366
1367 fn capture<'a>(sink: &'a CapturingSink) -> impl FnMut(&[u8]) -> std::io::Result<()> + 'a {
1368 move |bytes: &[u8]| {
1369 sink.chunks.borrow_mut().push(bytes.to_vec());
1370 Ok(())
1371 }
1372 }
1373
1374 #[test]
1375 fn chunk_producer_flushes_on_byte_cap() {
1376 let clock = FakeClock::new(0);
1377 let cfg = StreamConfig {
1378 chunk_default_pages: 1, chunk_min_pages: 1,
1380 chunk_max_pages: 1,
1381 chunk_max_rows: 1_000_000,
1382 chunk_max_latency_ms: 1_000_000,
1383 ..StreamConfig::DEFAULT
1384 };
1385 let sink = CapturingSink::new();
1386 let mut producer = ChunkProducer::new(&cfg, &clock);
1387 let mut flush = capture(&sink);
1388
1389 producer
1390 .push_line(&vec![b'x'; 8 * 1024], &mut flush)
1391 .unwrap();
1392 assert_eq!(sink.len(), 0);
1393
1394 let triggered = producer
1395 .push_line(&vec![b'y'; 8 * 1024], &mut flush)
1396 .unwrap();
1397 assert!(triggered);
1398 assert_eq!(sink.len(), 1);
1399 assert_eq!(sink.last_len(), Some(16 * 1024));
1400 assert_eq!(producer.last_flush_reason(), Some(FlushReason::Byte));
1401 }
1402
1403 #[test]
1404 fn chunk_producer_flushes_on_row_cap() {
1405 let clock = FakeClock::new(0);
1406 let cfg = StreamConfig {
1407 chunk_default_pages: 4, chunk_min_pages: 1,
1409 chunk_max_pages: 64,
1410 chunk_max_rows: 3,
1411 chunk_max_latency_ms: 1_000_000,
1412 ..StreamConfig::DEFAULT
1413 };
1414 let sink = CapturingSink::new();
1415 let mut producer = ChunkProducer::new(&cfg, &clock);
1416 let mut flush = capture(&sink);
1417 let row = b"{\"row\":{}}\n";
1418 producer.push_line(row, &mut flush).unwrap();
1419 producer.push_line(row, &mut flush).unwrap();
1420 assert_eq!(sink.len(), 0);
1421 let triggered = producer.push_line(row, &mut flush).unwrap();
1422 assert!(triggered);
1423 assert_eq!(sink.len(), 1);
1424 assert_eq!(producer.last_flush_reason(), Some(FlushReason::Row));
1425 }
1426
1427 #[test]
1428 fn chunk_producer_flushes_on_latency_cap() {
1429 let clock = FakeClock::new(0);
1430 let cfg = StreamConfig {
1431 chunk_default_pages: 4,
1432 chunk_min_pages: 1,
1433 chunk_max_pages: 64,
1434 chunk_max_rows: 1_000_000,
1435 chunk_max_latency_ms: 50,
1436 ..StreamConfig::DEFAULT
1437 };
1438 let sink = CapturingSink::new();
1439 let mut producer = ChunkProducer::new(&cfg, &clock);
1440 let mut flush = capture(&sink);
1441 producer.push_line(b"{\"row\":{}}\n", &mut flush).unwrap();
1442 assert_eq!(sink.len(), 0);
1443 clock.advance(60);
1444 let triggered = producer.push_line(b"{\"row\":{}}\n", &mut flush).unwrap();
1445 assert!(triggered);
1446 assert_eq!(producer.last_flush_reason(), Some(FlushReason::Latency));
1447 }
1448
1449 #[test]
1450 fn chunk_producer_finish_emits_terminal_flush() {
1451 let clock = FakeClock::new(0);
1452 let cfg = StreamConfig::DEFAULT;
1453 let sink = CapturingSink::new();
1454 let mut producer = ChunkProducer::new(&cfg, &clock);
1455 let mut flush = capture(&sink);
1456 producer.push_line(b"{\"row\":{}}\n", &mut flush).unwrap();
1457 producer.finish(&mut flush).unwrap();
1458 assert_eq!(sink.len(), 1);
1459 assert_eq!(producer.last_flush_reason(), Some(FlushReason::Terminal));
1460 }
1461
1462 #[test]
1463 fn write_chunked_helpers_produce_well_formed_chunks() {
1464 let mut buf: Vec<u8> = Vec::new();
1465 write_chunked_response_header(&mut buf, 200, "application/x-ndjson").unwrap();
1466 write_chunk(&mut buf, b"{\"row\":{}}\n").unwrap();
1467 write_chunked_terminator(&mut buf).unwrap();
1468 let text = String::from_utf8(buf).unwrap();
1469 assert!(text.starts_with("HTTP/1.1 200 OK\r\n"));
1470 assert!(text.contains("Transfer-Encoding: chunked\r\n"));
1471 assert!(text.contains("\r\nb\r\n{\"row\":{}}\n\r\n"));
1473 assert!(text.ends_with("0\r\n\r\n"));
1474 }
1475
1476 #[test]
1479 fn capacity_registry_global_exhausted_returns_structured_error() {
1480 let reg = StreamCapacityRegistry::new();
1481 let _g1 = reg.try_acquire("alice", 2, 32).unwrap();
1482 let _g2 = reg.try_acquire("alice", 2, 32).unwrap();
1483 let err = reg.try_acquire("alice", 2, 32).unwrap_err();
1484 assert_eq!(
1485 err,
1486 AcquireError::GlobalExhausted {
1487 limit: 2,
1488 current: 2,
1489 }
1490 );
1491 assert_eq!(err.code(), "server_stream_capacity_exhausted");
1492 }
1493
1494 #[test]
1495 fn capacity_registry_per_principal_exhausted_independent_of_global() {
1496 let reg = StreamCapacityRegistry::new();
1500 let _a1 = reg.try_acquire("alice", 100, 2).unwrap();
1501 let _a2 = reg.try_acquire("alice", 100, 2).unwrap();
1502 let err = reg.try_acquire("alice", 100, 2).unwrap_err();
1503 assert_eq!(
1504 err,
1505 AcquireError::PrincipalExhausted {
1506 principal: "alice".to_string(),
1507 limit: 2,
1508 current: 2,
1509 }
1510 );
1511 assert_eq!(err.code(), "principal_stream_quota_exhausted");
1512
1513 let _b1 = reg.try_acquire("bob", 100, 2).unwrap();
1515 let _b2 = reg.try_acquire("bob", 100, 2).unwrap();
1516 }
1517
1518 #[test]
1519 fn capacity_registry_release_frees_both_counters() {
1520 let reg = StreamCapacityRegistry::new();
1522 let g1 = reg.try_acquire("alice", 1, 1).unwrap();
1523 assert!(reg.try_acquire("alice", 1, 1).is_err());
1524 drop(g1);
1525 let (global, per_principal) = reg.snapshot();
1526 assert_eq!(global, 0);
1527 assert!(per_principal.is_empty());
1528 let _g2 = reg.try_acquire("alice", 1, 1).unwrap();
1530 }
1531
1532 #[test]
1533 fn capacity_registry_concurrent_acquire_release_does_not_over_issue() {
1534 use std::sync::atomic::{AtomicUsize, Ordering};
1540
1541 const THREADS: usize = 16;
1542 const ITERS: usize = 200;
1543 const CAP_GLOBAL: usize = 4;
1544 const CAP_PER_PRINCIPAL: usize = 4;
1545
1546 let reg = StreamCapacityRegistry::new();
1547 let observed_max = Arc::new(AtomicUsize::new(0));
1548 let mut handles = Vec::new();
1549 for tid in 0..THREADS {
1550 let reg = Arc::clone(®);
1551 let observed_max = Arc::clone(&observed_max);
1552 let principal = format!("p{}", tid % 2);
1555 handles.push(std::thread::spawn(move || {
1556 for _ in 0..ITERS {
1557 if let Ok(guard) = reg.try_acquire(&principal, CAP_GLOBAL, CAP_PER_PRINCIPAL) {
1558 let (live, _) = reg.snapshot();
1559 observed_max.fetch_max(live, Ordering::SeqCst);
1560 std::thread::yield_now();
1563 drop(guard);
1564 }
1565 }
1566 }));
1567 }
1568 for h in handles {
1569 h.join().unwrap();
1570 }
1571 let (global_after, per_principal_after) = reg.snapshot();
1572 assert_eq!(global_after, 0, "global counter leaked");
1573 assert!(
1574 per_principal_after.is_empty(),
1575 "per-principal map leaked: {per_principal_after:?}"
1576 );
1577 assert!(
1578 observed_max.load(Ordering::SeqCst) <= CAP_GLOBAL,
1579 "global cap was breached: observed {} > {}",
1580 observed_max.load(Ordering::SeqCst),
1581 CAP_GLOBAL
1582 );
1583 }
1584
1585 #[test]
1588 fn assess_resumability_accepts_plain_select() {
1589 assert!(assess_resumability("SELECT id, name FROM t"));
1590 assert!(assess_resumability("select * from t where id > 5"));
1591 assert!(assess_resumability("SELECT a, b FROM t ORDER BY rid"));
1592 assert!(assess_resumability("SELECT a, b FROM t ORDER BY rid ASC"));
1593 assert!(assess_resumability(
1594 "SELECT a FROM t ORDER BY rid ASC LIMIT 10"
1595 ));
1596 }
1597
1598 #[test]
1599 fn assess_resumability_rejects_aggregates_and_unordered() {
1600 assert!(!assess_resumability("SELECT COUNT(*) FROM t"));
1601 assert!(!assess_resumability("SELECT SUM(x) FROM t"));
1602 assert!(!assess_resumability("SELECT a, COUNT(b) FROM t GROUP BY a"));
1603 assert!(!assess_resumability("SELECT DISTINCT a FROM t"));
1604 assert!(!assess_resumability("SELECT a FROM t ORDER BY name"));
1605 assert!(!assess_resumability("SELECT a FROM t ORDER BY rid DESC"));
1606 assert!(!assess_resumability("SELECT a FROM t ORDER BY a, b"));
1607 assert!(!assess_resumability("INSERT INTO t (a) VALUES (1)"));
1608 assert!(!assess_resumability(
1609 "SELECT a FROM t JOIN u ON t.id = u.id"
1610 ));
1611 }
1612
1613 #[test]
1614 fn lease_registry_records_and_expires_against_ttl() {
1615 let reg = LeaseRegistry::new();
1616 reg.record(42, 1_000, 5_000);
1617 assert_eq!(reg.lookup(42, 1_000), LeaseLookup::Live);
1618 assert_eq!(reg.lookup(42, 5_999), LeaseLookup::Live);
1619 assert_eq!(reg.lookup(42, 6_000), LeaseLookup::Expired);
1620 assert_eq!(reg.lookup(99, 1_000), LeaseLookup::Unknown);
1621 }
1622
1623 #[test]
1624 fn prefix_hasher_is_order_sensitive_and_deterministic() {
1625 let mut a = PrefixHasher::new();
1626 a.update(b"{\"row\":{\"id\":1}}");
1627 a.update(b"{\"row\":{\"id\":2}}");
1628 let hash_a = a.finalize_hex();
1629
1630 let mut b = PrefixHasher::new();
1631 b.update(b"{\"row\":{\"id\":1}}");
1632 b.update(b"{\"row\":{\"id\":2}}");
1633 let hash_b = b.finalize_hex();
1634 assert_eq!(hash_a, hash_b);
1635
1636 let mut c = PrefixHasher::new();
1637 c.update(b"{\"row\":{\"id\":2}}");
1638 c.update(b"{\"row\":{\"id\":1}}");
1639 assert_ne!(hash_a, c.finalize_hex());
1640 assert_eq!(hash_a.len(), 64);
1641 }
1642
1643 #[test]
1646 fn lease_handle_is_128_bit_hex_and_unique_per_open() {
1647 let clock = FakeClock::new(0);
1651 let a = open_stream(StreamConfig::DEFAULT, 1, false, &clock).unwrap();
1652 let b = open_stream(StreamConfig::DEFAULT, 1, false, &clock).unwrap();
1653 assert_eq!(
1654 a.lease_handle.len(),
1655 LEASE_HANDLE_BYTES * 2,
1656 "handle must be 128 bits hex-encoded: {}",
1657 a.lease_handle
1658 );
1659 assert!(
1660 a.lease_handle.chars().all(|c| c.is_ascii_hexdigit()),
1661 "handle must be hex: {}",
1662 a.lease_handle
1663 );
1664 assert_ne!(a.lease_handle, b.lease_handle, "handles must differ");
1665 assert!(b.id > a.id);
1668 }
1669
1670 #[test]
1671 fn generate_lease_handle_produces_high_entropy_distinct_values() {
1672 let mut seen = std::collections::HashSet::new();
1676 for _ in 0..1024 {
1677 assert!(
1678 seen.insert(generate_lease_handle()),
1679 "duplicate handle in CSPRNG sequence"
1680 );
1681 }
1682 }
1683
1684 #[test]
1685 fn parse_jwt_exp_ms_extracts_seconds_to_ms() {
1686 let token = "eyJhbGciOiJIUzI1NiJ9.eyJleHAiOjE3MDAwMDAwMDB9.sig";
1691 assert_eq!(parse_jwt_exp_ms(token), Some(1_700_000_000_000));
1692 }
1693
1694 #[test]
1695 fn parse_jwt_exp_ms_returns_none_for_opaque_tokens() {
1696 assert_eq!(parse_jwt_exp_ms("not-a-jwt"), None);
1697 assert_eq!(parse_jwt_exp_ms("only.two"), None);
1698 assert_eq!(parse_jwt_exp_ms("a.b.c"), None);
1699 }
1700
1701 fn register_default_cursor(reg: &CursorRegistry, now_ms: u64) -> String {
1704 reg.register(
1705 42,
1706 "acme",
1707 "bearer:abc",
1708 "SELECT id FROM users ORDER BY rid",
1709 None,
1710 None,
1711 now_ms,
1712 1_000,
1713 )
1714 }
1715
1716 #[test]
1717 fn cursor_token_is_192_bit_hex_and_unique_per_register() {
1718 let reg = CursorRegistry::default();
1719 let a = register_default_cursor(®, 0);
1720 let b = register_default_cursor(®, 0);
1721 assert_eq!(
1722 a.len(),
1723 CURSOR_TOKEN_BYTES * 2,
1724 "token must be 192 bits hex-encoded: {a}"
1725 );
1726 assert!(
1727 a.chars().all(|c| c.is_ascii_hexdigit()),
1728 "token must be hex: {a}"
1729 );
1730 assert_ne!(a, b, "tokens must differ across registrations");
1731 assert_eq!(reg.len(), 2);
1732 }
1733
1734 #[test]
1735 fn cursor_resolves_for_owner_within_ttl() {
1736 let reg = CursorRegistry::default();
1738 let token = register_default_cursor(®, 1_000);
1739 let resume = reg
1740 .resolve(&token, "acme", "bearer:abc", 1_500)
1741 .expect("live cursor resolves for its owner");
1742 assert_eq!(resume.snapshot_lsn, 42, "resume re-pins the same snapshot");
1743 assert_eq!(resume.query, "SELECT id FROM users ORDER BY rid");
1744 assert_eq!(resume.expires_at_ms, 2_000);
1745 }
1746
1747 #[test]
1748 fn cursor_rejects_after_ttl_for_owner() {
1749 let reg = CursorRegistry::default();
1752 let token = register_default_cursor(®, 0);
1753 assert_eq!(
1754 reg.resolve(&token, "acme", "bearer:abc", 1_000),
1755 Err(CursorReject::Expired),
1756 "TTL boundary is inclusive — cursor is dead at opened_at + ttl"
1757 );
1758 assert_eq!(
1759 reg.resolve(&token, "acme", "bearer:abc", 5_000),
1760 Err(CursorReject::Expired)
1761 );
1762 }
1763
1764 #[test]
1765 fn cursor_cross_tenant_is_masked_as_not_found() {
1766 let reg = CursorRegistry::default();
1770 let token = register_default_cursor(®, 0);
1771 assert_eq!(
1772 reg.resolve(&token, "evil-corp", "bearer:abc", 100),
1773 Err(CursorReject::NotFound),
1774 "cross-tenant resume must mask existence as NotFound"
1775 );
1776 assert!(reg.resolve(&token, "acme", "bearer:abc", 100).is_ok());
1778 }
1779
1780 #[test]
1781 fn cursor_cross_principal_is_masked_as_not_found() {
1782 let reg = CursorRegistry::default();
1785 let token = register_default_cursor(®, 0);
1786 assert_eq!(
1787 reg.resolve(&token, "acme", "bearer:other", 100),
1788 Err(CursorReject::NotFound),
1789 "cross-principal resume must mask existence as NotFound"
1790 );
1791 assert!(reg.resolve(&token, "acme", "bearer:abc", 100).is_ok());
1792 }
1793
1794 #[test]
1795 fn cursor_unknown_token_is_not_found() {
1796 let reg = CursorRegistry::default();
1797 assert!(reg.is_empty());
1798 assert_eq!(
1799 reg.resolve("deadbeef", "acme", "bearer:abc", 0),
1800 Err(CursorReject::NotFound)
1801 );
1802 }
1803
1804 #[test]
1805 fn cursor_scope_is_checked_before_expiry() {
1806 let reg = CursorRegistry::default();
1810 let token = register_default_cursor(®, 0);
1811 assert_eq!(
1812 reg.resolve(&token, "evil-corp", "bearer:abc", 10_000),
1813 Err(CursorReject::NotFound),
1814 "expired + wrong scope must mask as NotFound, not Expired"
1815 );
1816 }
1817
1818 #[test]
1821 fn cancel_tombstones_cursor_and_raises_token() {
1822 let reg = CursorRegistry::default();
1825 let token = register_default_cursor(®, 0);
1826 let live = reg
1827 .cancel_token_for(&token)
1828 .expect("freshly-minted cursor exposes its token");
1829 assert!(!live.is_cancelled(), "token starts un-cancelled");
1830
1831 let returned = reg
1832 .cancel(&token, "acme", "bearer:abc")
1833 .expect("owner cancels its own cursor");
1834 assert!(returned.is_cancelled(), "cancel raises the returned token");
1835 assert!(
1836 live.is_cancelled(),
1837 "the handler-held clone observes the cancel"
1838 );
1839 }
1840
1841 #[test]
1842 fn cancelled_cursor_rejects_resume_with_cancelled_reason() {
1843 let reg = CursorRegistry::default();
1847 let token = register_default_cursor(®, 0);
1848 reg.cancel(&token, "acme", "bearer:abc")
1849 .expect("owner cancels");
1850 assert_eq!(
1851 reg.resolve(&token, "acme", "bearer:abc", 100),
1852 Err(CursorReject::Cancelled),
1853 "owner resuming a cancelled cursor sees Cancelled"
1854 );
1855 }
1856
1857 #[test]
1858 fn cancel_is_idempotent() {
1859 let reg = CursorRegistry::default();
1860 let token = register_default_cursor(®, 0);
1861 assert!(reg.cancel(&token, "acme", "bearer:abc").is_ok());
1862 let second = reg
1864 .cancel(&token, "acme", "bearer:abc")
1865 .expect("re-cancel is a no-op success");
1866 assert!(second.is_cancelled());
1867 }
1868
1869 #[test]
1870 fn cancel_cross_tenant_is_masked_as_not_found() {
1871 let reg = CursorRegistry::default();
1874 let token = register_default_cursor(®, 0);
1875 assert!(
1876 matches!(
1877 reg.cancel(&token, "evil-corp", "bearer:abc"),
1878 Err(CursorReject::NotFound)
1879 ),
1880 "cross-tenant cancel must mask existence as NotFound"
1881 );
1882 assert!(reg.resolve(&token, "acme", "bearer:abc", 100).is_ok());
1884 }
1885
1886 #[test]
1887 fn cancel_cross_principal_is_masked_as_not_found() {
1888 let reg = CursorRegistry::default();
1889 let token = register_default_cursor(®, 0);
1890 assert!(
1891 matches!(
1892 reg.cancel(&token, "acme", "bearer:other"),
1893 Err(CursorReject::NotFound)
1894 ),
1895 "cross-principal cancel must mask existence as NotFound"
1896 );
1897 assert!(reg.resolve(&token, "acme", "bearer:abc", 100).is_ok());
1898 }
1899
1900 #[test]
1901 fn cancel_unknown_token_is_not_found() {
1902 let reg = CursorRegistry::default();
1903 assert!(matches!(
1904 reg.cancel("deadbeef", "acme", "bearer:abc"),
1905 Err(CursorReject::NotFound)
1906 ));
1907 }
1908
1909 #[test]
1910 fn generate_cursor_token_produces_high_entropy_distinct_values() {
1911 let mut seen = std::collections::HashSet::new();
1912 for _ in 0..1024 {
1913 assert!(
1914 seen.insert(generate_cursor_token()),
1915 "duplicate cursor token in CSPRNG sequence"
1916 );
1917 }
1918 }
1919
1920 #[test]
1921 fn close_reason_as_str_covers_every_state_transition() {
1922 for (variant, expected) in [
1925 (CloseReason::Ok, "ok"),
1926 (CloseReason::Cancelled, "cancelled"),
1927 (CloseReason::Error, "error"),
1928 (CloseReason::SnapshotExpired, "snapshot_expired"),
1929 (CloseReason::CapacityRefused, "capacity_refused"),
1930 (CloseReason::IntegrityFailed, "integrity_failed"),
1931 ] {
1932 assert_eq!(variant.as_str(), expected);
1933 }
1934 }
1935
1936 #[test]
1937 fn stream_config_defaults_carry_s2_caps() {
1938 assert_eq!(StreamConfig::DEFAULT.max_global_streams, 256);
1939 assert_eq!(StreamConfig::DEFAULT.max_per_principal_streams, 32);
1940 }
1941
1942 struct SizeSink {
1947 sizes: std::cell::RefCell<Vec<usize>>,
1948 }
1949 impl SizeSink {
1950 fn new() -> Self {
1951 Self {
1952 sizes: std::cell::RefCell::new(Vec::new()),
1953 }
1954 }
1955 fn flushes(&self) -> usize {
1956 self.sizes.borrow().len()
1957 }
1958 fn max_chunk(&self) -> usize {
1959 self.sizes.borrow().iter().copied().max().unwrap_or(0)
1960 }
1961 }
1962 fn size_capture<'a>(sink: &'a SizeSink) -> impl FnMut(&[u8]) -> std::io::Result<()> + 'a {
1963 move |bytes: &[u8]| {
1964 sink.sizes.borrow_mut().push(bytes.len());
1965 Ok(())
1966 }
1967 }
1968
1969 #[test]
1970 fn drive_lines_streams_large_source_with_bounded_working_set() {
1971 let clock = FakeClock::new(0);
1978 let cfg = StreamConfig {
1979 chunk_default_pages: 1, chunk_min_pages: 1,
1981 chunk_max_pages: 1,
1982 chunk_max_rows: 1_000_000, chunk_max_latency_ms: 1_000_000,
1984 ..StreamConfig::DEFAULT
1985 };
1986 let sink = SizeSink::new();
1987 let mut producer = ChunkProducer::new(&cfg, &clock);
1988 let mut flush = size_capture(&sink);
1989
1990 const N: u64 = 1_000_000;
1991 let source = 0..N;
1993 let consumed = producer
1994 .drive_lines(
1995 source,
1996 |i: &u64| format!("{{\"row\":{{\"id\":{i}}}}}\n").into_bytes(),
1997 &mut flush,
1998 )
1999 .unwrap();
2000 producer.finish(&mut flush).unwrap();
2001
2002 assert_eq!(consumed, N);
2003 assert_eq!(producer.total_rows(), N);
2004 assert!(
2006 sink.flushes() > 1000,
2007 "expected the source to stream across many chunks, saw {}",
2008 sink.flushes()
2009 );
2010 let max_line = format!("{{\"row\":{{\"id\":{}}}}}\n", N - 1).len();
2013 assert!(
2014 sink.max_chunk() <= cfg.production_buffer_bytes() + max_line,
2015 "chunk {} exceeded bounded working set {}",
2016 sink.max_chunk(),
2017 cfg.production_buffer_bytes() + max_line
2018 );
2019 }
2020
2021 #[test]
2022 fn drive_lines_first_chunk_flushes_on_latency_before_source_drains() {
2023 let clock = FakeClock::new(0);
2029 let cfg = StreamConfig {
2030 chunk_default_pages: 64, chunk_min_pages: 1,
2032 chunk_max_pages: 64,
2033 chunk_max_rows: 1_000_000, chunk_max_latency_ms: 50,
2035 ..StreamConfig::DEFAULT
2036 };
2037 let sink = SizeSink::new();
2038 let mut producer = ChunkProducer::new(&cfg, &clock);
2039
2040 let mut first_flush_after: Option<u64> = None;
2044 let mut row = 0u64;
2045 while row < 1_000_000 {
2046 let line = format!("{{\"row\":{{\"id\":{row}}}}}\n");
2047 clock.advance(20);
2048 let mut flush = size_capture(&sink);
2049 let flushed = producer.push_line(line.as_bytes(), &mut flush).unwrap();
2050 row += 1;
2051 if flushed {
2052 first_flush_after = Some(row);
2053 break;
2054 }
2055 }
2056
2057 assert_eq!(producer.last_flush_reason(), Some(FlushReason::Latency));
2058 let rows_before_flush = first_flush_after.expect("a latency flush must occur");
2059 assert!(
2060 rows_before_flush <= 4,
2061 "first chunk flushed only after {rows_before_flush} rows; latency bound not honoured"
2062 );
2063 assert!(rows_before_flush < 1_000_000);
2065 }
2066
2067 #[test]
2068 fn drive_lines_parity_with_manual_push_line() {
2069 let clock = FakeClock::new(0);
2072 let cfg = StreamConfig::DEFAULT;
2073
2074 let lines: Vec<Vec<u8>> = (0..50)
2075 .map(|i| format!("{{\"row\":{{\"id\":{i}}}}}\n").into_bytes())
2076 .collect();
2077
2078 let driven = CapturingSink::new();
2079 {
2080 let mut p = ChunkProducer::new(&cfg, &clock);
2081 let mut flush = capture(&driven);
2082 p.drive_lines(lines.iter().cloned(), |l: &Vec<u8>| l.clone(), &mut flush)
2083 .unwrap();
2084 p.finish(&mut flush).unwrap();
2085 }
2086
2087 let manual = CapturingSink::new();
2088 {
2089 let mut p = ChunkProducer::new(&cfg, &clock);
2090 let mut flush = capture(&manual);
2091 for l in &lines {
2092 p.push_line(l, &mut flush).unwrap();
2093 }
2094 p.finish(&mut flush).unwrap();
2095 }
2096
2097 assert_eq!(*driven.chunks.borrow(), *manual.chunks.borrow());
2098 }
2099}