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
31pub const PAGE_SIZE: usize = 16 * 1024;
34
35pub trait Clock: Send + Sync {
38 fn now_ms(&self) -> u64;
39}
40
41#[derive(Debug, Default)]
42pub struct SystemClock;
43
44impl Clock for SystemClock {
45 fn now_ms(&self) -> u64 {
46 SystemTime::now()
47 .duration_since(UNIX_EPOCH)
48 .map(|d| d.as_millis() as u64)
49 .unwrap_or(0)
50 }
51}
52
53#[derive(Debug)]
54pub struct FakeClock {
55 now_ms: AtomicU64,
56}
57
58impl FakeClock {
59 pub fn new(now_ms: u64) -> Self {
60 Self {
61 now_ms: AtomicU64::new(now_ms),
62 }
63 }
64
65 pub fn advance(&self, ms: u64) {
66 self.now_ms.fetch_add(ms, Ordering::SeqCst);
67 }
68}
69
70impl Clock for FakeClock {
71 fn now_ms(&self) -> u64 {
72 self.now_ms.load(Ordering::SeqCst)
73 }
74}
75
76#[derive(Debug, Clone, Copy)]
81pub struct StreamConfig {
82 pub snapshot_ttl_ms: u64,
83 pub chunk_default_pages: usize,
84 pub chunk_min_pages: usize,
85 pub chunk_max_pages: usize,
86 pub chunk_max_rows: usize,
87 pub chunk_max_latency_ms: u64,
88 pub max_global_streams: usize,
92 pub max_per_principal_streams: usize,
97 pub default_verify: crate::runtime::integrity_tombstone::VerifyMode,
101}
102
103impl Default for StreamConfig {
104 fn default() -> Self {
105 Self::DEFAULT
106 }
107}
108
109impl StreamConfig {
110 pub const DEFAULT: StreamConfig = StreamConfig {
113 snapshot_ttl_ms: 60_000,
114 chunk_default_pages: 4,
115 chunk_min_pages: 1,
116 chunk_max_pages: 64,
117 chunk_max_rows: 1000,
118 chunk_max_latency_ms: 50,
119 max_global_streams: 256,
120 max_per_principal_streams: 32,
121 default_verify: crate::runtime::integrity_tombstone::VerifyMode::None,
122 };
123
124 pub fn load(runtime: &RedDBRuntime) -> crate::RedDBResult<Self> {
128 let store = runtime.db().store();
129 let read_u64 = |key: &str| -> crate::RedDBResult<Option<u64>> {
130 let invalid =
131 || crate::RedDBError::Query(format!("invalid {key}: expected an unsigned integer"));
132 match store.try_get_config(key)? {
133 None => Ok(None),
134 Some(Value::Integer(value)) => {
135 u64::try_from(value).map(Some).map_err(|_| invalid())
136 }
137 Some(Value::UnsignedInteger(value)) => Ok(Some(value)),
138 Some(Value::Text(text)) => text.parse().map(Some).map_err(|_| invalid()),
139 Some(_) => Err(invalid()),
140 }
141 };
142
143 let mut cfg = Self::DEFAULT;
144 if let Some(v) = read_u64("stream.snapshot.ttl_ms")? {
145 cfg.snapshot_ttl_ms = v;
146 }
147 if let Some(v) = read_u64("stream.chunk.default_pages")? {
148 cfg.chunk_default_pages = usize::try_from(v).map_err(|_| {
149 crate::RedDBError::Query(
150 "stream.chunk.default_pages exceeds platform size".to_string(),
151 )
152 })?;
153 }
154 if let Some(v) = read_u64("stream.chunk.min_pages")? {
155 cfg.chunk_min_pages = usize::try_from(v).map_err(|_| {
156 crate::RedDBError::Query("stream.chunk.min_pages exceeds platform size".to_string())
157 })?;
158 }
159 if let Some(v) = read_u64("stream.chunk.max_pages")? {
160 cfg.chunk_max_pages = usize::try_from(v).map_err(|_| {
161 crate::RedDBError::Query("stream.chunk.max_pages exceeds platform size".to_string())
162 })?;
163 }
164 if let Some(v) = read_u64("stream.chunk.max_rows")? {
165 cfg.chunk_max_rows = usize::try_from(v).map_err(|_| {
166 crate::RedDBError::Query("stream.chunk.max_rows exceeds platform size".to_string())
167 })?;
168 }
169 if let Some(v) = read_u64("stream.chunk.max_latency_ms")? {
170 cfg.chunk_max_latency_ms = v;
171 }
172 if let Some(v) = read_u64("stream.max_global")? {
173 cfg.max_global_streams = usize::try_from(v).map_err(|_| {
174 crate::RedDBError::Query("stream.max_global exceeds platform size".to_string())
175 })?;
176 }
177 if let Some(v) = read_u64("stream.max_per_principal")? {
178 cfg.max_per_principal_streams = usize::try_from(v).map_err(|_| {
179 crate::RedDBError::Query(
180 "stream.max_per_principal exceeds platform size".to_string(),
181 )
182 })?;
183 }
184 match store.try_get_config("stream.integrity.default_verify")? {
185 None => {}
186 Some(Value::Text(text))
187 if matches!(text.trim().to_ascii_lowercase().as_str(), "none" | "sha256") =>
188 {
189 cfg.default_verify = crate::runtime::integrity_tombstone::VerifyMode::parse(&text);
190 }
191 Some(_) => {
192 return Err(crate::RedDBError::Query(
193 "invalid stream.integrity.default_verify: expected none or sha256".to_string(),
194 ))
195 }
196 }
197 cfg.normalize();
198 Ok(cfg)
199 }
200
201 fn normalize(&mut self) {
205 if self.chunk_min_pages == 0 {
206 self.chunk_min_pages = 1;
207 }
208 if self.chunk_max_pages < self.chunk_min_pages {
209 self.chunk_max_pages = self.chunk_min_pages;
210 }
211 if self.chunk_default_pages < self.chunk_min_pages {
212 self.chunk_default_pages = self.chunk_min_pages;
213 }
214 if self.chunk_default_pages > self.chunk_max_pages {
215 self.chunk_default_pages = self.chunk_max_pages;
216 }
217 if self.chunk_max_rows == 0 {
218 self.chunk_max_rows = 1;
219 }
220 if self.max_global_streams == 0 {
221 self.max_global_streams = 1;
222 }
223 if self.max_per_principal_streams == 0 {
224 self.max_per_principal_streams = 1;
225 }
226 }
227
228 pub fn production_buffer_bytes(&self) -> usize {
231 self.chunk_default_pages.saturating_mul(PAGE_SIZE)
232 }
233}
234
235static NEXT_LEASE_ID: AtomicU64 = AtomicU64::new(1);
243
244pub const LEASE_HANDLE_BYTES: usize = 16;
250
251pub fn generate_lease_handle() -> String {
255 let mut bytes = [0u8; LEASE_HANDLE_BYTES];
256 if crate::crypto::os_random::fill_bytes(&mut bytes).is_err() {
257 let lo = NEXT_LEASE_ID.fetch_add(0, Ordering::SeqCst).to_le_bytes();
263 let now = crate::utils::now_unix_nanos().to_le_bytes();
264 bytes[..8].copy_from_slice(&lo);
265 bytes[8..].copy_from_slice(&now);
266 }
267 crate::utils::to_hex(&bytes)
268}
269
270pub fn parse_jwt_exp_ms(token: &str) -> Option<u64> {
280 let parts: Vec<&str> = token.split('.').collect();
281 if parts.len() != 3 {
282 return None;
283 }
284 let payload = base64url_decode_padded(parts[1])?;
285 let v: crate::json::Value = crate::json::from_slice(&payload).ok()?;
286 let exp_secs = v.get("exp").and_then(|n| n.as_f64())? as i64;
287 if exp_secs <= 0 {
288 return None;
289 }
290 Some((exp_secs as u64).saturating_mul(1000))
291}
292
293fn base64url_decode_padded(input: &str) -> Option<Vec<u8>> {
294 let mut s = String::with_capacity(input.len() + 4);
295 for ch in input.chars() {
296 match ch {
297 '-' => s.push('+'),
298 '_' => s.push('/'),
299 _ => s.push(ch),
300 }
301 }
302 while !s.len().is_multiple_of(4) {
303 s.push('=');
304 }
305 reddb_wire::redwire::handshake::base64_std_decode(&s)
306}
307
308#[derive(Debug, Clone, Copy, PartialEq, Eq)]
312pub enum CloseReason {
313 Ok,
314 Cancelled,
315 Error,
316 SnapshotExpired,
317 CapacityRefused,
318 IntegrityFailed,
319}
320
321impl CloseReason {
322 pub fn as_str(&self) -> &'static str {
323 match self {
324 CloseReason::Ok => "ok",
325 CloseReason::Cancelled => "cancelled",
326 CloseReason::Error => "error",
327 CloseReason::SnapshotExpired => "snapshot_expired",
328 CloseReason::CapacityRefused => "capacity_refused",
329 CloseReason::IntegrityFailed => "integrity_failed",
330 }
331 }
332}
333
334#[derive(Debug)]
335pub struct StreamLease {
336 pub id: u64,
337 pub lease_handle: String,
341 pub snapshot_lsn: u64,
342 pub opened_at_ms: u64,
343 pub config: StreamConfig,
344}
345
346impl StreamLease {
347 pub fn snapshot_expired(&self, now_ms: u64) -> bool {
354 now_ms.saturating_sub(self.opened_at_ms) >= self.config.snapshot_ttl_ms
355 }
356}
357
358#[derive(Debug, PartialEq, Eq)]
359pub enum OpenStreamError {
360 TransactionActive,
364}
365
366impl OpenStreamError {
367 pub fn code(&self) -> &'static str {
368 match self {
369 OpenStreamError::TransactionActive => "stream_in_transaction_unsupported",
370 }
371 }
372
373 pub fn message(&self) -> &'static str {
374 match self {
375 OpenStreamError::TransactionActive => {
376 "cannot open output stream while a transaction is active on this session"
377 }
378 }
379 }
380}
381
382pub fn open_stream(
387 config: StreamConfig,
388 snapshot_lsn: u64,
389 in_transaction: bool,
390 clock: &dyn Clock,
391) -> Result<StreamLease, OpenStreamError> {
392 if in_transaction {
393 return Err(OpenStreamError::TransactionActive);
394 }
395 Ok(StreamLease {
396 id: NEXT_LEASE_ID.fetch_add(1, Ordering::SeqCst),
397 lease_handle: generate_lease_handle(),
398 snapshot_lsn,
399 opened_at_ms: clock.now_ms(),
400 config,
401 })
402}
403
404pub struct ChunkProducer<'a> {
412 buf: Vec<u8>,
413 rows: usize,
414 window_started_ms: u64,
415 cap_bytes: usize,
416 cap_rows: usize,
417 cap_latency_ms: u64,
418 clock: &'a dyn Clock,
419 total_flushes: u64,
420 total_bytes: u64,
421 total_rows: u64,
422 last_flush_reason: Option<FlushReason>,
423}
424
425#[derive(Debug, Clone, Copy, PartialEq, Eq)]
426pub enum FlushReason {
427 Byte,
428 Row,
429 Latency,
430 Terminal,
431}
432
433impl<'a> ChunkProducer<'a> {
434 pub fn new(config: &StreamConfig, clock: &'a dyn Clock) -> Self {
435 let cap_bytes = config.production_buffer_bytes();
436 Self {
437 buf: Vec::with_capacity(cap_bytes),
438 rows: 0,
439 window_started_ms: clock.now_ms(),
440 cap_bytes,
441 cap_rows: config.chunk_max_rows,
442 cap_latency_ms: config.chunk_max_latency_ms,
443 clock,
444 total_flushes: 0,
445 total_bytes: 0,
446 total_rows: 0,
447 last_flush_reason: None,
448 }
449 }
450
451 pub fn push_line<F>(&mut self, line: &[u8], flush: &mut F) -> std::io::Result<bool>
454 where
455 F: FnMut(&[u8]) -> std::io::Result<()>,
456 {
457 crate::runtime::query_budget::charge_output(line.len()).map_err(std::io::Error::other)?;
458 self.buf.extend_from_slice(line);
459 self.rows += 1;
460 self.total_rows += 1;
461
462 if self.buf.len() >= self.cap_bytes {
463 self.flush(flush, FlushReason::Byte)?;
464 return Ok(true);
465 }
466 if self.rows >= self.cap_rows {
467 self.flush(flush, FlushReason::Row)?;
468 return Ok(true);
469 }
470 let elapsed = self.clock.now_ms().saturating_sub(self.window_started_ms);
471 if elapsed >= self.cap_latency_ms {
472 self.flush(flush, FlushReason::Latency)?;
473 return Ok(true);
474 }
475 Ok(false)
476 }
477
478 pub fn drive_lines<S, R, Enc, F>(
496 &mut self,
497 source: S,
498 mut encode: Enc,
499 flush: &mut F,
500 ) -> std::io::Result<u64>
501 where
502 S: IntoIterator<Item = R>,
503 Enc: FnMut(&R) -> Vec<u8>,
504 F: FnMut(&[u8]) -> std::io::Result<()>,
505 {
506 let mut count = 0u64;
507 for record in source {
508 let line = encode(&record);
509 self.push_line(&line, flush)?;
510 count += 1;
511 }
512 Ok(count)
513 }
514
515 pub fn finish<F>(&mut self, flush: &mut F) -> std::io::Result<()>
519 where
520 F: FnMut(&[u8]) -> std::io::Result<()>,
521 {
522 if !self.buf.is_empty() {
523 self.flush(flush, FlushReason::Terminal)?;
524 }
525 Ok(())
526 }
527
528 fn flush<F>(&mut self, flush: &mut F, reason: FlushReason) -> std::io::Result<()>
529 where
530 F: FnMut(&[u8]) -> std::io::Result<()>,
531 {
532 flush(&self.buf)?;
533 self.total_bytes += self.buf.len() as u64;
534 self.total_flushes += 1;
535 self.last_flush_reason = Some(reason);
536 self.buf.clear();
537 self.rows = 0;
538 self.window_started_ms = self.clock.now_ms();
539 Ok(())
540 }
541
542 pub fn total_flushes(&self) -> u64 {
543 self.total_flushes
544 }
545 pub fn total_bytes(&self) -> u64 {
546 self.total_bytes
547 }
548 pub fn total_rows(&self) -> u64 {
549 self.total_rows
550 }
551 pub fn last_flush_reason(&self) -> Option<FlushReason> {
552 self.last_flush_reason
553 }
554}
555
556pub fn write_chunked_response_header<W: std::io::Write>(
560 writer: &mut W,
561 status: u16,
562 content_type: &str,
563) -> std::io::Result<()> {
564 let mut header = format!(
572 "HTTP/1.1 {} {}\r\nContent-Type: {}\r\nTransfer-Encoding: chunked\r\nCache-Control: no-cache\r\nConnection: close\r\n",
573 status,
574 crate::server::transport::status_text(status),
575 content_type,
576 );
577 for (name, value) in crate::server::transport::cors_header_pairs() {
581 header.push_str(name);
582 header.push_str(": ");
583 header.push_str(&value);
584 header.push_str("\r\n");
585 }
586 header.push_str("\r\n");
587 writer.write_all(header.as_bytes())?;
588 writer.flush()
589}
590
591pub fn write_chunk<W: std::io::Write>(writer: &mut W, bytes: &[u8]) -> std::io::Result<()> {
595 if bytes.is_empty() {
596 return Ok(());
597 }
598 let size = format!("{:x}\r\n", bytes.len());
599 writer.write_all(size.as_bytes())?;
600 writer.write_all(bytes)?;
601 writer.write_all(b"\r\n")?;
602 writer.flush()
603}
604
605pub fn write_chunked_terminator<W: std::io::Write>(writer: &mut W) -> std::io::Result<()> {
607 writer.write_all(b"0\r\n\r\n")?;
608 writer.flush()
609}
610
611#[derive(Debug, Default)]
619pub struct StreamCapacityRegistry {
620 inner: Mutex<CapacityInner>,
621}
622
623#[derive(Debug, Default)]
624struct CapacityInner {
625 global_count: usize,
626 per_principal: HashMap<String, usize>,
627}
628
629#[derive(Debug, Clone, PartialEq, Eq)]
634pub enum AcquireError {
635 GlobalExhausted { limit: usize, current: usize },
637 PrincipalExhausted {
641 principal: String,
642 limit: usize,
643 current: usize,
644 },
645}
646
647impl AcquireError {
648 pub fn code(&self) -> &'static str {
649 match self {
650 AcquireError::GlobalExhausted { .. } => "server_stream_capacity_exhausted",
651 AcquireError::PrincipalExhausted { .. } => "principal_stream_quota_exhausted",
652 }
653 }
654
655 pub fn message(&self) -> String {
656 match self {
657 AcquireError::GlobalExhausted { limit, current } => {
658 format!("server stream capacity exhausted (limit {limit}, current {current})")
659 }
660 AcquireError::PrincipalExhausted {
661 principal,
662 limit,
663 current,
664 } => format!(
665 "principal {principal} stream quota exhausted (limit {limit}, current {current})"
666 ),
667 }
668 }
669}
670
671impl StreamCapacityRegistry {
672 pub fn new() -> Arc<Self> {
673 Arc::new(Self::default())
674 }
675
676 pub fn try_acquire(
680 self: &Arc<Self>,
681 principal: &str,
682 max_global: usize,
683 max_per_principal: usize,
684 ) -> Result<StreamCapacityGuard, AcquireError> {
685 let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
686 if inner.global_count >= max_global {
687 return Err(AcquireError::GlobalExhausted {
688 limit: max_global,
689 current: inner.global_count,
690 });
691 }
692 let current = inner.per_principal.get(principal).copied().unwrap_or(0);
693 if current >= max_per_principal {
694 return Err(AcquireError::PrincipalExhausted {
695 principal: principal.to_string(),
696 limit: max_per_principal,
697 current,
698 });
699 }
700 inner.global_count += 1;
701 inner
702 .per_principal
703 .insert(principal.to_string(), current + 1);
704 Ok(StreamCapacityGuard {
705 registry: Arc::clone(self),
706 principal: principal.to_string(),
707 released: false,
708 })
709 }
710
711 fn release(&self, principal: &str) {
712 let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
713 if inner.global_count > 0 {
714 inner.global_count -= 1;
715 }
716 if let Some(count) = inner.per_principal.get_mut(principal) {
717 if *count > 0 {
718 *count -= 1;
719 }
720 if *count == 0 {
721 inner.per_principal.remove(principal);
722 }
723 }
724 }
725
726 pub fn snapshot(&self) -> (usize, HashMap<String, usize>) {
728 let inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
729 (inner.global_count, inner.per_principal.clone())
730 }
731}
732
733#[must_use = "dropping the guard immediately releases the stream slot"]
738#[derive(Debug)]
739pub struct StreamCapacityGuard {
740 registry: Arc<StreamCapacityRegistry>,
741 principal: String,
742 released: bool,
743}
744
745impl StreamCapacityGuard {
746 pub fn principal(&self) -> &str {
747 &self.principal
748 }
749}
750
751impl Drop for StreamCapacityGuard {
752 fn drop(&mut self) {
753 if !self.released {
754 self.registry.release(&self.principal);
755 self.released = true;
756 }
757 }
758}
759
760pub fn assess_resumability(query: &str) -> bool {
769 let upper = query.to_uppercase();
770 let trimmed = upper.trim_start();
771 if !trimmed.starts_with("SELECT ") && !trimmed.starts_with("SELECT\n") {
772 return false;
773 }
774 const FORBIDDEN: &[&str] = &[
775 " GROUP BY ",
776 " HAVING ",
777 " DISTINCT ",
778 "DISTINCT ",
779 "COUNT(",
780 "SUM(",
781 "AVG(",
782 "MIN(",
783 "MAX(",
784 "ARRAY_AGG(",
785 "JSON_AGG(",
786 "OVER(",
787 " OVER (",
788 " JOIN ",
789 ];
790 for kw in FORBIDDEN {
791 if upper.contains(kw) {
792 return false;
793 }
794 }
795 if let Some(idx) = upper.find("ORDER BY") {
796 let tail = &upper[idx + "ORDER BY".len()..];
797 let mut clause = tail.to_string();
799 if let Some(lim) = clause.find(" LIMIT ") {
800 clause.truncate(lim);
801 }
802 if let Some(semi) = clause.find(';') {
803 clause.truncate(semi);
804 }
805 let clause = clause.trim();
806 if !matches!(clause, "RID" | "RID ASC") {
807 return false;
808 }
809 }
810 true
811}
812
813#[derive(Debug, Default)]
819pub struct LeaseRegistry {
820 inner: Mutex<HashMap<u64, LeaseEntry>>,
821}
822
823#[derive(Debug, Clone, Copy)]
824struct LeaseEntry {
825 opened_at_ms: u64,
826 ttl_ms: u64,
827}
828
829#[derive(Debug, Clone, Copy, PartialEq, Eq)]
830pub enum LeaseLookup {
831 Unknown,
833 Expired,
835 Live,
837}
838
839impl LeaseRegistry {
840 pub fn new() -> Arc<Self> {
841 Arc::new(Self::default())
842 }
843
844 pub fn record(&self, snapshot_lsn: u64, opened_at_ms: u64, ttl_ms: u64) {
849 let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
850 inner.insert(
851 snapshot_lsn,
852 LeaseEntry {
853 opened_at_ms,
854 ttl_ms,
855 },
856 );
857 }
858
859 pub fn lookup(&self, snapshot_lsn: u64, now_ms: u64) -> LeaseLookup {
862 let inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
863 match inner.get(&snapshot_lsn) {
864 None => LeaseLookup::Unknown,
865 Some(entry) => {
866 if now_ms.saturating_sub(entry.opened_at_ms) >= entry.ttl_ms {
867 LeaseLookup::Expired
868 } else {
869 LeaseLookup::Live
870 }
871 }
872 }
873 }
874
875 #[doc(hidden)]
877 pub fn len(&self) -> usize {
878 self.inner.lock().unwrap_or_else(|e| e.into_inner()).len()
879 }
880}
881
882pub const CURSOR_TOKEN_BYTES: usize = 24;
889
890static NEXT_CURSOR_SALT: AtomicU64 = AtomicU64::new(1);
891
892pub fn generate_cursor_token() -> String {
897 let mut bytes = [0u8; CURSOR_TOKEN_BYTES];
898 if crate::crypto::os_random::fill_bytes(&mut bytes).is_err() {
899 let salt = NEXT_CURSOR_SALT
900 .fetch_add(1, Ordering::SeqCst)
901 .to_le_bytes();
902 let now = crate::utils::now_unix_nanos().to_le_bytes();
903 bytes[..8].copy_from_slice(&salt);
904 bytes[8..16].copy_from_slice(&now);
905 }
906 crate::utils::to_hex(&bytes)
907}
908
909#[derive(Debug, Clone)]
913struct CursorEntry {
914 snapshot_lsn: u64,
915 tenant: String,
916 principal: String,
917 query: String,
918 entity_types: Option<Vec<String>>,
919 capabilities: Option<Vec<String>>,
920 opened_at_ms: u64,
921 ttl_ms: u64,
922 cancel_token: CancelToken,
927 cancelled: bool,
931}
932
933#[derive(Debug, Clone, PartialEq, Eq)]
937pub struct CursorResume {
938 pub snapshot_lsn: u64,
939 pub query: String,
940 pub entity_types: Option<Vec<String>>,
941 pub capabilities: Option<Vec<String>>,
942 pub expires_at_ms: u64,
943}
944
945#[derive(Debug, Clone, Copy, PartialEq, Eq)]
950pub enum CursorReject {
951 NotFound,
954 Expired,
956 Cancelled,
961}
962
963#[derive(Debug, Default)]
974pub struct CursorRegistry {
975 inner: Mutex<HashMap<String, CursorEntry>>,
976}
977
978impl CursorRegistry {
979 pub fn new() -> Arc<Self> {
980 Arc::new(Self::default())
981 }
982
983 #[allow(clippy::too_many_arguments)]
987 pub fn register(
988 &self,
989 snapshot_lsn: u64,
990 tenant: &str,
991 principal: &str,
992 query: &str,
993 entity_types: Option<Vec<String>>,
994 capabilities: Option<Vec<String>>,
995 opened_at_ms: u64,
996 ttl_ms: u64,
997 ) -> String {
998 let token = generate_cursor_token();
999 let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
1000 inner.insert(
1001 token.clone(),
1002 CursorEntry {
1003 snapshot_lsn,
1004 tenant: tenant.to_string(),
1005 principal: principal.to_string(),
1006 query: query.to_string(),
1007 entity_types,
1008 capabilities,
1009 opened_at_ms,
1010 ttl_ms,
1011 cancel_token: CancelToken::new(),
1012 cancelled: false,
1013 },
1014 );
1015 token
1016 }
1017
1018 pub fn cancel_token_for(&self, token: &str) -> Option<CancelToken> {
1023 let inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
1024 inner.get(token).map(|e| e.cancel_token.clone())
1025 }
1026
1027 pub fn cancel(
1038 &self,
1039 token: &str,
1040 tenant: &str,
1041 principal: &str,
1042 ) -> Result<CancelToken, CursorReject> {
1043 let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
1044 let entry = inner.get_mut(token).ok_or(CursorReject::NotFound)?;
1045 if entry.tenant != tenant || entry.principal != principal {
1046 return Err(CursorReject::NotFound);
1047 }
1048 entry.cancelled = true;
1049 entry.cancel_token.cancel();
1050 Ok(entry.cancel_token.clone())
1051 }
1052
1053 pub fn resolve(
1061 &self,
1062 token: &str,
1063 tenant: &str,
1064 principal: &str,
1065 now_ms: u64,
1066 ) -> Result<CursorResume, CursorReject> {
1067 let inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
1068 let entry = inner.get(token).ok_or(CursorReject::NotFound)?;
1069 if entry.tenant != tenant || entry.principal != principal {
1073 return Err(CursorReject::NotFound);
1074 }
1075 if entry.cancelled {
1079 return Err(CursorReject::Cancelled);
1080 }
1081 if now_ms.saturating_sub(entry.opened_at_ms) >= entry.ttl_ms {
1082 return Err(CursorReject::Expired);
1083 }
1084 Ok(CursorResume {
1085 snapshot_lsn: entry.snapshot_lsn,
1086 query: entry.query.clone(),
1087 entity_types: entry.entity_types.clone(),
1088 capabilities: entry.capabilities.clone(),
1089 expires_at_ms: entry.opened_at_ms.saturating_add(entry.ttl_ms),
1090 })
1091 }
1092
1093 #[doc(hidden)]
1095 pub fn len(&self) -> usize {
1096 self.inner.lock().unwrap_or_else(|e| e.into_inner()).len()
1097 }
1098
1099 #[doc(hidden)]
1100 pub fn is_empty(&self) -> bool {
1101 self.len() == 0
1102 }
1103}
1104
1105#[derive(Debug, Default)]
1111pub struct PrefixHasher {
1112 inner: Option<sha2::Sha256>,
1113 rows: u64,
1114}
1115
1116impl PrefixHasher {
1117 pub fn new() -> Self {
1118 use sha2::Digest;
1119 Self {
1120 inner: Some(sha2::Sha256::new()),
1121 rows: 0,
1122 }
1123 }
1124
1125 pub fn update(&mut self, line: &[u8]) {
1126 use sha2::Digest;
1127 if let Some(h) = self.inner.as_mut() {
1128 h.update(line);
1129 }
1130 self.rows += 1;
1131 }
1132
1133 pub fn rows(&self) -> u64 {
1134 self.rows
1135 }
1136
1137 pub fn finalize_hex(mut self) -> String {
1140 use sha2::Digest;
1141 let hasher = self
1142 .inner
1143 .take()
1144 .expect("PrefixHasher::finalize_hex called twice");
1145 let digest = hasher.finalize();
1146 let mut out = String::with_capacity(64);
1147 for b in digest.iter() {
1148 out.push_str(&format!("{b:02x}"));
1149 }
1150 out
1151 }
1152}
1153
1154pub fn audit_stream_opened(
1160 runtime: &RedDBRuntime,
1161 lease_handle: &str,
1162 principal: &str,
1163 snapshot_lsn: u64,
1164 query_hash: &str,
1165) {
1166 use crate::json::{Map, Value as JsonValue};
1167 let mut detail = Map::new();
1168 detail.insert(
1169 "lease_handle".to_string(),
1170 JsonValue::String(lease_handle.to_string()),
1171 );
1172 detail.insert(
1173 "snapshot_lsn".to_string(),
1174 JsonValue::Number(snapshot_lsn as f64),
1175 );
1176 detail.insert(
1177 "query_hash".to_string(),
1178 JsonValue::String(query_hash.to_string()),
1179 );
1180 let event = crate::runtime::audit_log::AuditEvent::builder("stream.opened")
1181 .principal(principal)
1182 .resource(lease_handle.to_string())
1183 .outcome(crate::runtime::audit_log::Outcome::Success)
1184 .detail(JsonValue::Object(detail))
1185 .build();
1186 runtime.audit_log().record_event(event);
1187}
1188
1189pub fn audit_stream_closed(
1190 runtime: &RedDBRuntime,
1191 lease_handle: &str,
1192 principal: &str,
1193 reason: CloseReason,
1194 row_count: u64,
1195 bytes_written: u64,
1196) {
1197 use crate::json::{Map, Value as JsonValue};
1198 let mut stats = Map::new();
1199 stats.insert("row_count".to_string(), JsonValue::Number(row_count as f64));
1200 stats.insert(
1201 "bytes_written".to_string(),
1202 JsonValue::Number(bytes_written as f64),
1203 );
1204 let mut detail = Map::new();
1205 detail.insert(
1206 "lease_handle".to_string(),
1207 JsonValue::String(lease_handle.to_string()),
1208 );
1209 detail.insert(
1210 "reason".to_string(),
1211 JsonValue::String(reason.as_str().to_string()),
1212 );
1213 detail.insert("stats".to_string(), JsonValue::Object(stats));
1214 let outcome = match reason {
1215 CloseReason::Ok => crate::runtime::audit_log::Outcome::Success,
1216 CloseReason::CapacityRefused => crate::runtime::audit_log::Outcome::Denied,
1217 _ => crate::runtime::audit_log::Outcome::Error,
1218 };
1219 let event = crate::runtime::audit_log::AuditEvent::builder("stream.closed")
1220 .principal(principal)
1221 .resource(lease_handle.to_string())
1222 .outcome(outcome)
1223 .detail(JsonValue::Object(detail))
1224 .build();
1225 runtime.audit_log().record_event(event);
1226}
1227
1228pub fn audit_token_expired_during_lease(
1229 runtime: &RedDBRuntime,
1230 lease_handle: &str,
1231 principal: &str,
1232 token_expiry_ms: u64,
1233) {
1234 use crate::json::{Map, Value as JsonValue};
1235 let mut detail = Map::new();
1236 detail.insert(
1237 "lease_handle".to_string(),
1238 JsonValue::String(lease_handle.to_string()),
1239 );
1240 detail.insert(
1241 "token_expiry".to_string(),
1242 JsonValue::Number(token_expiry_ms as f64),
1243 );
1244 detail.insert("lease_continued".to_string(), JsonValue::Bool(true));
1245 let event = crate::runtime::audit_log::AuditEvent::builder("stream.token_expired_during_lease")
1246 .principal(principal)
1247 .resource(lease_handle.to_string())
1248 .outcome(crate::runtime::audit_log::Outcome::Success)
1249 .detail(JsonValue::Object(detail))
1250 .build();
1251 runtime.audit_log().record_event(event);
1252}
1253
1254pub fn audit_stream_capacity_refused(
1260 runtime: &RedDBRuntime,
1261 principal: &str,
1262 code: &str,
1263 limit: usize,
1264 current: usize,
1265) {
1266 use crate::json::{Map, Value as JsonValue};
1267 let mut detail = Map::new();
1268 detail.insert(
1269 "reason".to_string(),
1270 JsonValue::String(CloseReason::CapacityRefused.as_str().to_string()),
1271 );
1272 detail.insert("code".to_string(), JsonValue::String(code.to_string()));
1273 detail.insert("limit".to_string(), JsonValue::Number(limit as f64));
1274 detail.insert("current".to_string(), JsonValue::Number(current as f64));
1275 let event = crate::runtime::audit_log::AuditEvent::builder("stream.closed")
1276 .principal(principal)
1277 .outcome(crate::runtime::audit_log::Outcome::Denied)
1278 .detail(JsonValue::Object(detail))
1279 .build();
1280 runtime.audit_log().record_event(event);
1281}
1282
1283pub fn system_clock() -> Arc<dyn Clock> {
1286 static INSTANCE: std::sync::OnceLock<Arc<dyn Clock>> = std::sync::OnceLock::new();
1287 Arc::clone(INSTANCE.get_or_init(|| Arc::new(SystemClock)))
1288}
1289
1290#[cfg(test)]
1291mod tests {
1292 use super::*;
1293
1294 #[test]
1295 fn open_stream_refuses_when_session_has_active_transaction() {
1296 let clock = FakeClock::new(0);
1297 let err = open_stream(StreamConfig::DEFAULT, 42, true, &clock).unwrap_err();
1298 assert_eq!(err, OpenStreamError::TransactionActive);
1299 assert_eq!(err.code(), "stream_in_transaction_unsupported");
1300 }
1301
1302 #[test]
1303 fn open_stream_succeeds_when_session_is_autocommit() {
1304 let clock = FakeClock::new(1_700_000_000_000);
1305 let lease = open_stream(StreamConfig::DEFAULT, 123, false, &clock).unwrap();
1306 assert_eq!(lease.snapshot_lsn, 123);
1307 assert_eq!(lease.opened_at_ms, 1_700_000_000_000);
1308 assert!(lease.id >= 1);
1309 }
1310
1311 #[test]
1312 fn lease_ids_are_unique_and_monotonic() {
1313 let clock = FakeClock::new(0);
1314 let a = open_stream(StreamConfig::DEFAULT, 1, false, &clock).unwrap();
1315 let b = open_stream(StreamConfig::DEFAULT, 1, false, &clock).unwrap();
1316 assert!(b.id > a.id);
1317 }
1318
1319 #[test]
1320 fn snapshot_expired_uses_injected_clock_and_ttl() {
1321 let clock = FakeClock::new(0);
1324 let mut config = StreamConfig::DEFAULT;
1325 config.snapshot_ttl_ms = 5_000;
1326 let lease = open_stream(config, 0, false, &clock).unwrap();
1327
1328 assert!(!lease.snapshot_expired(clock.now_ms()));
1329 clock.advance(4_999);
1330 assert!(!lease.snapshot_expired(clock.now_ms()));
1331 clock.advance(1);
1332 assert!(lease.snapshot_expired(clock.now_ms()));
1333 }
1334
1335 #[test]
1336 fn stream_config_loads_defaults_when_kv_is_empty() {
1337 let cfg = StreamConfig::DEFAULT;
1339 assert_eq!(cfg.snapshot_ttl_ms, 60_000);
1340 assert_eq!(cfg.chunk_default_pages, 4);
1341 assert_eq!(cfg.chunk_min_pages, 1);
1342 assert_eq!(cfg.chunk_max_pages, 64);
1343 assert_eq!(cfg.chunk_max_rows, 1000);
1344 assert_eq!(cfg.chunk_max_latency_ms, 50);
1345 assert_eq!(cfg.production_buffer_bytes(), 64 * 1024);
1346 }
1347
1348 #[test]
1349 fn stream_config_normalize_clamps_inconsistent_inputs() {
1350 let mut cfg = StreamConfig {
1351 snapshot_ttl_ms: 1,
1352 chunk_default_pages: 100,
1353 chunk_min_pages: 0,
1354 chunk_max_pages: 8,
1355 chunk_max_rows: 0,
1356 chunk_max_latency_ms: 1,
1357 max_global_streams: 0,
1358 max_per_principal_streams: 0,
1359 default_verify: crate::runtime::integrity_tombstone::VerifyMode::None,
1360 };
1361 cfg.normalize();
1362 assert_eq!(cfg.chunk_min_pages, 1);
1363 assert_eq!(cfg.chunk_max_pages, 8);
1364 assert_eq!(cfg.chunk_default_pages, 8); assert!(cfg.chunk_max_rows >= 1);
1366 assert!(cfg.max_global_streams >= 1);
1367 assert!(cfg.max_per_principal_streams >= 1);
1368 }
1369
1370 struct CapturingSink {
1375 chunks: std::cell::RefCell<Vec<Vec<u8>>>,
1376 }
1377 impl CapturingSink {
1378 fn new() -> Self {
1379 Self {
1380 chunks: std::cell::RefCell::new(Vec::new()),
1381 }
1382 }
1383 fn len(&self) -> usize {
1384 self.chunks.borrow().len()
1385 }
1386 fn last_len(&self) -> Option<usize> {
1387 self.chunks.borrow().last().map(|c| c.len())
1388 }
1389 }
1390
1391 fn capture<'a>(sink: &'a CapturingSink) -> impl FnMut(&[u8]) -> std::io::Result<()> + 'a {
1392 move |bytes: &[u8]| {
1393 sink.chunks.borrow_mut().push(bytes.to_vec());
1394 Ok(())
1395 }
1396 }
1397
1398 #[test]
1399 fn chunk_producer_flushes_on_byte_cap() {
1400 let clock = FakeClock::new(0);
1401 let cfg = StreamConfig {
1402 chunk_default_pages: 1, chunk_min_pages: 1,
1404 chunk_max_pages: 1,
1405 chunk_max_rows: 1_000_000,
1406 chunk_max_latency_ms: 1_000_000,
1407 ..StreamConfig::DEFAULT
1408 };
1409 let sink = CapturingSink::new();
1410 let mut producer = ChunkProducer::new(&cfg, &clock);
1411 let mut flush = capture(&sink);
1412
1413 producer
1414 .push_line(&vec![b'x'; 8 * 1024], &mut flush)
1415 .unwrap();
1416 assert_eq!(sink.len(), 0);
1417
1418 let triggered = producer
1419 .push_line(&vec![b'y'; 8 * 1024], &mut flush)
1420 .unwrap();
1421 assert!(triggered);
1422 assert_eq!(sink.len(), 1);
1423 assert_eq!(sink.last_len(), Some(16 * 1024));
1424 assert_eq!(producer.last_flush_reason(), Some(FlushReason::Byte));
1425 }
1426
1427 #[test]
1428 fn chunk_producer_flushes_on_row_cap() {
1429 let clock = FakeClock::new(0);
1430 let cfg = StreamConfig {
1431 chunk_default_pages: 4, chunk_min_pages: 1,
1433 chunk_max_pages: 64,
1434 chunk_max_rows: 3,
1435 chunk_max_latency_ms: 1_000_000,
1436 ..StreamConfig::DEFAULT
1437 };
1438 let sink = CapturingSink::new();
1439 let mut producer = ChunkProducer::new(&cfg, &clock);
1440 let mut flush = capture(&sink);
1441 let row = b"{\"row\":{}}\n";
1442 producer.push_line(row, &mut flush).unwrap();
1443 producer.push_line(row, &mut flush).unwrap();
1444 assert_eq!(sink.len(), 0);
1445 let triggered = producer.push_line(row, &mut flush).unwrap();
1446 assert!(triggered);
1447 assert_eq!(sink.len(), 1);
1448 assert_eq!(producer.last_flush_reason(), Some(FlushReason::Row));
1449 }
1450
1451 #[test]
1452 fn chunk_producer_flushes_on_latency_cap() {
1453 let clock = FakeClock::new(0);
1454 let cfg = StreamConfig {
1455 chunk_default_pages: 4,
1456 chunk_min_pages: 1,
1457 chunk_max_pages: 64,
1458 chunk_max_rows: 1_000_000,
1459 chunk_max_latency_ms: 50,
1460 ..StreamConfig::DEFAULT
1461 };
1462 let sink = CapturingSink::new();
1463 let mut producer = ChunkProducer::new(&cfg, &clock);
1464 let mut flush = capture(&sink);
1465 producer.push_line(b"{\"row\":{}}\n", &mut flush).unwrap();
1466 assert_eq!(sink.len(), 0);
1467 clock.advance(60);
1468 let triggered = producer.push_line(b"{\"row\":{}}\n", &mut flush).unwrap();
1469 assert!(triggered);
1470 assert_eq!(producer.last_flush_reason(), Some(FlushReason::Latency));
1471 }
1472
1473 #[test]
1474 fn chunk_producer_finish_emits_terminal_flush() {
1475 let clock = FakeClock::new(0);
1476 let cfg = StreamConfig::DEFAULT;
1477 let sink = CapturingSink::new();
1478 let mut producer = ChunkProducer::new(&cfg, &clock);
1479 let mut flush = capture(&sink);
1480 producer.push_line(b"{\"row\":{}}\n", &mut flush).unwrap();
1481 producer.finish(&mut flush).unwrap();
1482 assert_eq!(sink.len(), 1);
1483 assert_eq!(producer.last_flush_reason(), Some(FlushReason::Terminal));
1484 }
1485
1486 #[test]
1487 fn write_chunked_helpers_produce_well_formed_chunks() {
1488 let mut buf: Vec<u8> = Vec::new();
1489 write_chunked_response_header(&mut buf, 200, "application/x-ndjson").unwrap();
1490 write_chunk(&mut buf, b"{\"row\":{}}\n").unwrap();
1491 write_chunked_terminator(&mut buf).unwrap();
1492 let text = String::from_utf8(buf).unwrap();
1493 assert!(text.starts_with("HTTP/1.1 200 OK\r\n"));
1494 assert!(text.contains("Transfer-Encoding: chunked\r\n"));
1495 assert!(text.contains("\r\nb\r\n{\"row\":{}}\n\r\n"));
1497 assert!(text.ends_with("0\r\n\r\n"));
1498 }
1499
1500 #[test]
1503 fn capacity_registry_global_exhausted_returns_structured_error() {
1504 let reg = StreamCapacityRegistry::new();
1505 let _g1 = reg.try_acquire("alice", 2, 32).unwrap();
1506 let _g2 = reg.try_acquire("alice", 2, 32).unwrap();
1507 let err = reg.try_acquire("alice", 2, 32).unwrap_err();
1508 assert_eq!(
1509 err,
1510 AcquireError::GlobalExhausted {
1511 limit: 2,
1512 current: 2,
1513 }
1514 );
1515 assert_eq!(err.code(), "server_stream_capacity_exhausted");
1516 }
1517
1518 #[test]
1519 fn capacity_registry_per_principal_exhausted_independent_of_global() {
1520 let reg = StreamCapacityRegistry::new();
1524 let _a1 = reg.try_acquire("alice", 100, 2).unwrap();
1525 let _a2 = reg.try_acquire("alice", 100, 2).unwrap();
1526 let err = reg.try_acquire("alice", 100, 2).unwrap_err();
1527 assert_eq!(
1528 err,
1529 AcquireError::PrincipalExhausted {
1530 principal: "alice".to_string(),
1531 limit: 2,
1532 current: 2,
1533 }
1534 );
1535 assert_eq!(err.code(), "principal_stream_quota_exhausted");
1536
1537 let _b1 = reg.try_acquire("bob", 100, 2).unwrap();
1539 let _b2 = reg.try_acquire("bob", 100, 2).unwrap();
1540 }
1541
1542 #[test]
1543 fn capacity_registry_release_frees_both_counters() {
1544 let reg = StreamCapacityRegistry::new();
1546 let g1 = reg.try_acquire("alice", 1, 1).unwrap();
1547 assert!(reg.try_acquire("alice", 1, 1).is_err());
1548 drop(g1);
1549 let (global, per_principal) = reg.snapshot();
1550 assert_eq!(global, 0);
1551 assert!(per_principal.is_empty());
1552 let _g2 = reg.try_acquire("alice", 1, 1).unwrap();
1554 }
1555
1556 #[test]
1557 fn capacity_registry_concurrent_acquire_release_does_not_over_issue() {
1558 use std::sync::atomic::{AtomicUsize, Ordering};
1564
1565 const THREADS: usize = 16;
1566 const ITERS: usize = 200;
1567 const CAP_GLOBAL: usize = 4;
1568 const CAP_PER_PRINCIPAL: usize = 4;
1569
1570 let reg = StreamCapacityRegistry::new();
1571 let observed_max = Arc::new(AtomicUsize::new(0));
1572 let mut handles = Vec::new();
1573 for tid in 0..THREADS {
1574 let reg = Arc::clone(®);
1575 let observed_max = Arc::clone(&observed_max);
1576 let principal = format!("p{}", tid % 2);
1579 handles.push(std::thread::spawn(move || {
1580 for _ in 0..ITERS {
1581 if let Ok(guard) = reg.try_acquire(&principal, CAP_GLOBAL, CAP_PER_PRINCIPAL) {
1582 let (live, _) = reg.snapshot();
1583 observed_max.fetch_max(live, Ordering::SeqCst);
1584 std::thread::yield_now();
1587 drop(guard);
1588 }
1589 }
1590 }));
1591 }
1592 for h in handles {
1593 h.join().unwrap();
1594 }
1595 let (global_after, per_principal_after) = reg.snapshot();
1596 assert_eq!(global_after, 0, "global counter leaked");
1597 assert!(
1598 per_principal_after.is_empty(),
1599 "per-principal map leaked: {per_principal_after:?}"
1600 );
1601 assert!(
1602 observed_max.load(Ordering::SeqCst) <= CAP_GLOBAL,
1603 "global cap was breached: observed {} > {}",
1604 observed_max.load(Ordering::SeqCst),
1605 CAP_GLOBAL
1606 );
1607 }
1608
1609 #[test]
1612 fn assess_resumability_accepts_plain_select() {
1613 assert!(assess_resumability("SELECT id, name FROM t"));
1614 assert!(assess_resumability("select * from t where id > 5"));
1615 assert!(assess_resumability("SELECT a, b FROM t ORDER BY rid"));
1616 assert!(assess_resumability("SELECT a, b FROM t ORDER BY rid ASC"));
1617 assert!(assess_resumability(
1618 "SELECT a FROM t ORDER BY rid ASC LIMIT 10"
1619 ));
1620 }
1621
1622 #[test]
1623 fn assess_resumability_rejects_aggregates_and_unordered() {
1624 assert!(!assess_resumability("SELECT COUNT(*) FROM t"));
1625 assert!(!assess_resumability("SELECT SUM(x) FROM t"));
1626 assert!(!assess_resumability("SELECT a, COUNT(b) FROM t GROUP BY a"));
1627 assert!(!assess_resumability("SELECT DISTINCT a FROM t"));
1628 assert!(!assess_resumability("SELECT a FROM t ORDER BY name"));
1629 assert!(!assess_resumability("SELECT a FROM t ORDER BY rid DESC"));
1630 assert!(!assess_resumability("SELECT a FROM t ORDER BY a, b"));
1631 assert!(!assess_resumability("INSERT INTO t (a) VALUES (1)"));
1632 assert!(!assess_resumability(
1633 "SELECT a FROM t JOIN u ON t.id = u.id"
1634 ));
1635 }
1636
1637 #[test]
1638 fn lease_registry_records_and_expires_against_ttl() {
1639 let reg = LeaseRegistry::new();
1640 reg.record(42, 1_000, 5_000);
1641 assert_eq!(reg.lookup(42, 1_000), LeaseLookup::Live);
1642 assert_eq!(reg.lookup(42, 5_999), LeaseLookup::Live);
1643 assert_eq!(reg.lookup(42, 6_000), LeaseLookup::Expired);
1644 assert_eq!(reg.lookup(99, 1_000), LeaseLookup::Unknown);
1645 }
1646
1647 #[test]
1648 fn prefix_hasher_is_order_sensitive_and_deterministic() {
1649 let mut a = PrefixHasher::new();
1650 a.update(b"{\"row\":{\"id\":1}}");
1651 a.update(b"{\"row\":{\"id\":2}}");
1652 let hash_a = a.finalize_hex();
1653
1654 let mut b = PrefixHasher::new();
1655 b.update(b"{\"row\":{\"id\":1}}");
1656 b.update(b"{\"row\":{\"id\":2}}");
1657 let hash_b = b.finalize_hex();
1658 assert_eq!(hash_a, hash_b);
1659
1660 let mut c = PrefixHasher::new();
1661 c.update(b"{\"row\":{\"id\":2}}");
1662 c.update(b"{\"row\":{\"id\":1}}");
1663 assert_ne!(hash_a, c.finalize_hex());
1664 assert_eq!(hash_a.len(), 64);
1665 }
1666
1667 #[test]
1670 fn lease_handle_is_128_bit_hex_and_unique_per_open() {
1671 let clock = FakeClock::new(0);
1675 let a = open_stream(StreamConfig::DEFAULT, 1, false, &clock).unwrap();
1676 let b = open_stream(StreamConfig::DEFAULT, 1, false, &clock).unwrap();
1677 assert_eq!(
1678 a.lease_handle.len(),
1679 LEASE_HANDLE_BYTES * 2,
1680 "handle must be 128 bits hex-encoded: {}",
1681 a.lease_handle
1682 );
1683 assert!(
1684 a.lease_handle.chars().all(|c| c.is_ascii_hexdigit()),
1685 "handle must be hex: {}",
1686 a.lease_handle
1687 );
1688 assert_ne!(a.lease_handle, b.lease_handle, "handles must differ");
1689 assert!(b.id > a.id);
1692 }
1693
1694 #[test]
1695 fn generate_lease_handle_produces_high_entropy_distinct_values() {
1696 let mut seen = std::collections::HashSet::new();
1700 for _ in 0..1024 {
1701 assert!(
1702 seen.insert(generate_lease_handle()),
1703 "duplicate handle in CSPRNG sequence"
1704 );
1705 }
1706 }
1707
1708 #[test]
1709 fn parse_jwt_exp_ms_extracts_seconds_to_ms() {
1710 let token = "eyJhbGciOiJIUzI1NiJ9.eyJleHAiOjE3MDAwMDAwMDB9.sig";
1715 assert_eq!(parse_jwt_exp_ms(token), Some(1_700_000_000_000));
1716 }
1717
1718 #[test]
1719 fn parse_jwt_exp_ms_returns_none_for_opaque_tokens() {
1720 assert_eq!(parse_jwt_exp_ms("not-a-jwt"), None);
1721 assert_eq!(parse_jwt_exp_ms("only.two"), None);
1722 assert_eq!(parse_jwt_exp_ms("a.b.c"), None);
1723 }
1724
1725 fn register_default_cursor(reg: &CursorRegistry, now_ms: u64) -> String {
1728 reg.register(
1729 42,
1730 "acme",
1731 "bearer:abc",
1732 "SELECT id FROM users ORDER BY rid",
1733 None,
1734 None,
1735 now_ms,
1736 1_000,
1737 )
1738 }
1739
1740 #[test]
1741 fn cursor_token_is_192_bit_hex_and_unique_per_register() {
1742 let reg = CursorRegistry::default();
1743 let a = register_default_cursor(®, 0);
1744 let b = register_default_cursor(®, 0);
1745 assert_eq!(
1746 a.len(),
1747 CURSOR_TOKEN_BYTES * 2,
1748 "token must be 192 bits hex-encoded: {a}"
1749 );
1750 assert!(
1751 a.chars().all(|c| c.is_ascii_hexdigit()),
1752 "token must be hex: {a}"
1753 );
1754 assert_ne!(a, b, "tokens must differ across registrations");
1755 assert_eq!(reg.len(), 2);
1756 }
1757
1758 #[test]
1759 fn cursor_resolves_for_owner_within_ttl() {
1760 let reg = CursorRegistry::default();
1762 let token = register_default_cursor(®, 1_000);
1763 let resume = reg
1764 .resolve(&token, "acme", "bearer:abc", 1_500)
1765 .expect("live cursor resolves for its owner");
1766 assert_eq!(resume.snapshot_lsn, 42, "resume re-pins the same snapshot");
1767 assert_eq!(resume.query, "SELECT id FROM users ORDER BY rid");
1768 assert_eq!(resume.expires_at_ms, 2_000);
1769 }
1770
1771 #[test]
1772 fn cursor_rejects_after_ttl_for_owner() {
1773 let reg = CursorRegistry::default();
1776 let token = register_default_cursor(®, 0);
1777 assert_eq!(
1778 reg.resolve(&token, "acme", "bearer:abc", 1_000),
1779 Err(CursorReject::Expired),
1780 "TTL boundary is inclusive — cursor is dead at opened_at + ttl"
1781 );
1782 assert_eq!(
1783 reg.resolve(&token, "acme", "bearer:abc", 5_000),
1784 Err(CursorReject::Expired)
1785 );
1786 }
1787
1788 #[test]
1789 fn cursor_cross_tenant_is_masked_as_not_found() {
1790 let reg = CursorRegistry::default();
1794 let token = register_default_cursor(®, 0);
1795 assert_eq!(
1796 reg.resolve(&token, "evil-corp", "bearer:abc", 100),
1797 Err(CursorReject::NotFound),
1798 "cross-tenant resume must mask existence as NotFound"
1799 );
1800 assert!(reg.resolve(&token, "acme", "bearer:abc", 100).is_ok());
1802 }
1803
1804 #[test]
1805 fn cursor_cross_principal_is_masked_as_not_found() {
1806 let reg = CursorRegistry::default();
1809 let token = register_default_cursor(®, 0);
1810 assert_eq!(
1811 reg.resolve(&token, "acme", "bearer:other", 100),
1812 Err(CursorReject::NotFound),
1813 "cross-principal resume must mask existence as NotFound"
1814 );
1815 assert!(reg.resolve(&token, "acme", "bearer:abc", 100).is_ok());
1816 }
1817
1818 #[test]
1819 fn cursor_unknown_token_is_not_found() {
1820 let reg = CursorRegistry::default();
1821 assert!(reg.is_empty());
1822 assert_eq!(
1823 reg.resolve("deadbeef", "acme", "bearer:abc", 0),
1824 Err(CursorReject::NotFound)
1825 );
1826 }
1827
1828 #[test]
1829 fn cursor_scope_is_checked_before_expiry() {
1830 let reg = CursorRegistry::default();
1834 let token = register_default_cursor(®, 0);
1835 assert_eq!(
1836 reg.resolve(&token, "evil-corp", "bearer:abc", 10_000),
1837 Err(CursorReject::NotFound),
1838 "expired + wrong scope must mask as NotFound, not Expired"
1839 );
1840 }
1841
1842 #[test]
1845 fn cancel_tombstones_cursor_and_raises_token() {
1846 let reg = CursorRegistry::default();
1849 let token = register_default_cursor(®, 0);
1850 let live = reg
1851 .cancel_token_for(&token)
1852 .expect("freshly-minted cursor exposes its token");
1853 assert!(!live.is_cancelled(), "token starts un-cancelled");
1854
1855 let returned = reg
1856 .cancel(&token, "acme", "bearer:abc")
1857 .expect("owner cancels its own cursor");
1858 assert!(returned.is_cancelled(), "cancel raises the returned token");
1859 assert!(
1860 live.is_cancelled(),
1861 "the handler-held clone observes the cancel"
1862 );
1863 }
1864
1865 #[test]
1866 fn cancelled_cursor_rejects_resume_with_cancelled_reason() {
1867 let reg = CursorRegistry::default();
1871 let token = register_default_cursor(®, 0);
1872 reg.cancel(&token, "acme", "bearer:abc")
1873 .expect("owner cancels");
1874 assert_eq!(
1875 reg.resolve(&token, "acme", "bearer:abc", 100),
1876 Err(CursorReject::Cancelled),
1877 "owner resuming a cancelled cursor sees Cancelled"
1878 );
1879 }
1880
1881 #[test]
1882 fn cancel_is_idempotent() {
1883 let reg = CursorRegistry::default();
1884 let token = register_default_cursor(®, 0);
1885 assert!(reg.cancel(&token, "acme", "bearer:abc").is_ok());
1886 let second = reg
1888 .cancel(&token, "acme", "bearer:abc")
1889 .expect("re-cancel is a no-op success");
1890 assert!(second.is_cancelled());
1891 }
1892
1893 #[test]
1894 fn cancel_cross_tenant_is_masked_as_not_found() {
1895 let reg = CursorRegistry::default();
1898 let token = register_default_cursor(®, 0);
1899 assert!(
1900 matches!(
1901 reg.cancel(&token, "evil-corp", "bearer:abc"),
1902 Err(CursorReject::NotFound)
1903 ),
1904 "cross-tenant cancel must mask existence as NotFound"
1905 );
1906 assert!(reg.resolve(&token, "acme", "bearer:abc", 100).is_ok());
1908 }
1909
1910 #[test]
1911 fn cancel_cross_principal_is_masked_as_not_found() {
1912 let reg = CursorRegistry::default();
1913 let token = register_default_cursor(®, 0);
1914 assert!(
1915 matches!(
1916 reg.cancel(&token, "acme", "bearer:other"),
1917 Err(CursorReject::NotFound)
1918 ),
1919 "cross-principal cancel must mask existence as NotFound"
1920 );
1921 assert!(reg.resolve(&token, "acme", "bearer:abc", 100).is_ok());
1922 }
1923
1924 #[test]
1925 fn cancel_unknown_token_is_not_found() {
1926 let reg = CursorRegistry::default();
1927 assert!(matches!(
1928 reg.cancel("deadbeef", "acme", "bearer:abc"),
1929 Err(CursorReject::NotFound)
1930 ));
1931 }
1932
1933 #[test]
1934 fn generate_cursor_token_produces_high_entropy_distinct_values() {
1935 let mut seen = std::collections::HashSet::new();
1936 for _ in 0..1024 {
1937 assert!(
1938 seen.insert(generate_cursor_token()),
1939 "duplicate cursor token in CSPRNG sequence"
1940 );
1941 }
1942 }
1943
1944 #[test]
1945 fn close_reason_as_str_covers_every_state_transition() {
1946 for (variant, expected) in [
1949 (CloseReason::Ok, "ok"),
1950 (CloseReason::Cancelled, "cancelled"),
1951 (CloseReason::Error, "error"),
1952 (CloseReason::SnapshotExpired, "snapshot_expired"),
1953 (CloseReason::CapacityRefused, "capacity_refused"),
1954 (CloseReason::IntegrityFailed, "integrity_failed"),
1955 ] {
1956 assert_eq!(variant.as_str(), expected);
1957 }
1958 }
1959
1960 #[test]
1961 fn stream_config_defaults_carry_s2_caps() {
1962 assert_eq!(StreamConfig::DEFAULT.max_global_streams, 256);
1963 assert_eq!(StreamConfig::DEFAULT.max_per_principal_streams, 32);
1964 }
1965
1966 struct SizeSink {
1971 sizes: std::cell::RefCell<Vec<usize>>,
1972 }
1973 impl SizeSink {
1974 fn new() -> Self {
1975 Self {
1976 sizes: std::cell::RefCell::new(Vec::new()),
1977 }
1978 }
1979 fn flushes(&self) -> usize {
1980 self.sizes.borrow().len()
1981 }
1982 fn max_chunk(&self) -> usize {
1983 self.sizes.borrow().iter().copied().max().unwrap_or(0)
1984 }
1985 }
1986 fn size_capture<'a>(sink: &'a SizeSink) -> impl FnMut(&[u8]) -> std::io::Result<()> + 'a {
1987 move |bytes: &[u8]| {
1988 sink.sizes.borrow_mut().push(bytes.len());
1989 Ok(())
1990 }
1991 }
1992
1993 #[test]
1994 fn drive_lines_streams_large_source_with_bounded_working_set() {
1995 let clock = FakeClock::new(0);
2002 let cfg = StreamConfig {
2003 chunk_default_pages: 1, chunk_min_pages: 1,
2005 chunk_max_pages: 1,
2006 chunk_max_rows: 1_000_000, chunk_max_latency_ms: 1_000_000,
2008 ..StreamConfig::DEFAULT
2009 };
2010 let sink = SizeSink::new();
2011 let mut producer = ChunkProducer::new(&cfg, &clock);
2012 let mut flush = size_capture(&sink);
2013
2014 const N: u64 = 1_000_000;
2015 let source = 0..N;
2017 let consumed = producer
2018 .drive_lines(
2019 source,
2020 |i: &u64| format!("{{\"row\":{{\"id\":{i}}}}}\n").into_bytes(),
2021 &mut flush,
2022 )
2023 .unwrap();
2024 producer.finish(&mut flush).unwrap();
2025
2026 assert_eq!(consumed, N);
2027 assert_eq!(producer.total_rows(), N);
2028 assert!(
2030 sink.flushes() > 1000,
2031 "expected the source to stream across many chunks, saw {}",
2032 sink.flushes()
2033 );
2034 let max_line = format!("{{\"row\":{{\"id\":{}}}}}\n", N - 1).len();
2037 assert!(
2038 sink.max_chunk() <= cfg.production_buffer_bytes() + max_line,
2039 "chunk {} exceeded bounded working set {}",
2040 sink.max_chunk(),
2041 cfg.production_buffer_bytes() + max_line
2042 );
2043 }
2044
2045 #[test]
2046 fn drive_lines_first_chunk_flushes_on_latency_before_source_drains() {
2047 let clock = FakeClock::new(0);
2053 let cfg = StreamConfig {
2054 chunk_default_pages: 64, chunk_min_pages: 1,
2056 chunk_max_pages: 64,
2057 chunk_max_rows: 1_000_000, chunk_max_latency_ms: 50,
2059 ..StreamConfig::DEFAULT
2060 };
2061 let sink = SizeSink::new();
2062 let mut producer = ChunkProducer::new(&cfg, &clock);
2063
2064 let mut first_flush_after: Option<u64> = None;
2068 let mut row = 0u64;
2069 while row < 1_000_000 {
2070 let line = format!("{{\"row\":{{\"id\":{row}}}}}\n");
2071 clock.advance(20);
2072 let mut flush = size_capture(&sink);
2073 let flushed = producer.push_line(line.as_bytes(), &mut flush).unwrap();
2074 row += 1;
2075 if flushed {
2076 first_flush_after = Some(row);
2077 break;
2078 }
2079 }
2080
2081 assert_eq!(producer.last_flush_reason(), Some(FlushReason::Latency));
2082 let rows_before_flush = first_flush_after.expect("a latency flush must occur");
2083 assert!(
2084 rows_before_flush <= 4,
2085 "first chunk flushed only after {rows_before_flush} rows; latency bound not honoured"
2086 );
2087 assert!(rows_before_flush < 1_000_000);
2089 }
2090
2091 #[test]
2092 fn drive_lines_parity_with_manual_push_line() {
2093 let clock = FakeClock::new(0);
2096 let cfg = StreamConfig::DEFAULT;
2097
2098 let lines: Vec<Vec<u8>> = (0..50)
2099 .map(|i| format!("{{\"row\":{{\"id\":{i}}}}}\n").into_bytes())
2100 .collect();
2101
2102 let driven = CapturingSink::new();
2103 {
2104 let mut p = ChunkProducer::new(&cfg, &clock);
2105 let mut flush = capture(&driven);
2106 p.drive_lines(lines.iter().cloned(), |l: &Vec<u8>| l.clone(), &mut flush)
2107 .unwrap();
2108 p.finish(&mut flush).unwrap();
2109 }
2110
2111 let manual = CapturingSink::new();
2112 {
2113 let mut p = ChunkProducer::new(&cfg, &clock);
2114 let mut flush = capture(&manual);
2115 for l in &lines {
2116 p.push_line(l, &mut flush).unwrap();
2117 }
2118 p.finish(&mut flush).unwrap();
2119 }
2120
2121 assert_eq!(*driven.chunks.borrow(), *manual.chunks.borrow());
2122 }
2123}