use std::collections::HashSet;
use tracing::{debug, info_span, warn};
use super::statement_frame::ReadFrame;
use super::RedDBRuntime;
use crate::api::{RedDBError, RedDBResult};
use crate::application::SearchContextInput;
use crate::storage::unified::devx::SimilarResult;
pub enum AuthorizedSearch {}
impl AuthorizedSearch {
pub(crate) fn execute_similar(
runtime: &RedDBRuntime,
scope: &dyn ReadFrame,
collection: &str,
vector: &[f32],
k: usize,
min_score: f32,
) -> RedDBResult<Vec<SimilarResult>> {
let span = info_span!(
"authorized_search.similar",
collection = collection,
tenant = ?scope.effective_scope(),
);
let _enter = span.enter();
let visible = require_visible(scope, "SEARCH SIMILAR")?;
if !visible.contains(collection) {
warn!(
target: "authorized_search",
collection = collection,
"denied: collection outside visible scope"
);
return Err(RedDBError::Query(format!(
"permission denied: collection `{collection}` is not in the caller's visible scope"
)));
}
debug!(target: "authorized_search", "scope-checked, dispatching");
with_read_frame(runtime, scope, || {
runtime.search_similar(collection, vector, k, min_score)
})
}
#[allow(clippy::too_many_arguments)]
pub(crate) fn execute_text(
runtime: &RedDBRuntime,
scope: &dyn ReadFrame,
query: String,
collections: Option<Vec<String>>,
entity_types: Option<Vec<String>>,
capabilities: Option<Vec<String>>,
fields: Option<Vec<String>>,
limit: Option<usize>,
fuzzy: bool,
) -> RedDBResult<crate::storage::unified::dsl::QueryResult> {
let span = info_span!(
"authorized_search.text",
tenant = ?scope.effective_scope(),
);
let _enter = span.enter();
let visible = require_visible(scope, "SEARCH TEXT")?;
let constrained = constrain_collections(collections, visible);
if let Some(ref c) = constrained {
if c.is_empty() {
return Err(RedDBError::Query(
"permission denied: no requested collection is in the caller's visible scope"
.to_string(),
));
}
}
with_read_frame(runtime, scope, || {
runtime.search_text(
query,
constrained,
entity_types,
capabilities,
fields,
limit,
fuzzy,
)
})
}
pub(crate) fn execute_context(
runtime: &RedDBRuntime,
scope: &dyn ReadFrame,
mut input: SearchContextInput,
) -> RedDBResult<crate::runtime::ContextSearchResult> {
let span = info_span!(
"authorized_search.context",
tenant = ?scope.effective_scope(),
);
let _enter = span.enter();
let visible = require_visible(scope, "SEARCH CONTEXT")?;
input.collections = constrain_collections(input.collections, visible);
if let Some(ref c) = input.collections {
if c.is_empty() {
return Err(RedDBError::Query(
"permission denied: no requested collection is in the caller's visible scope"
.to_string(),
));
}
} else {
let mut bounded: Vec<String> = visible.iter().cloned().collect();
bounded.sort();
input.collections = Some(bounded);
}
let mut result = with_read_frame(runtime, scope, || runtime.search_context(input))?;
post_filter_context_result(&mut result, visible);
Ok(result)
}
}
fn with_read_frame<R>(runtime: &RedDBRuntime, frame: &dyn ReadFrame, run: impl FnOnce() -> R) -> R {
let bundle = super::execution_context::SnapshotBundle {
snapshot: Some(frame.snapshot_context(runtime)),
auth: frame
.identity()
.map(|(name, role)| (name.to_string(), role)),
tenant: frame.effective_scope().map(str::to_string),
};
super::execution_context::with_snapshot_bundle(&bundle, run)
}
fn post_filter_context_result(
result: &mut crate::runtime::ContextSearchResult,
allowed: &HashSet<String>,
) {
let retain = |bucket: &mut Vec<crate::runtime::ContextEntity>| {
bucket.retain(|e| allowed.contains(&e.collection));
};
retain(&mut result.tables);
retain(&mut result.graph.nodes);
retain(&mut result.graph.edges);
retain(&mut result.vectors);
retain(&mut result.documents);
retain(&mut result.key_values);
let visible_ids: HashSet<u64> = std::iter::empty()
.chain(result.tables.iter().map(|e| e.entity.id.raw()))
.chain(result.graph.nodes.iter().map(|e| e.entity.id.raw()))
.chain(result.graph.edges.iter().map(|e| e.entity.id.raw()))
.chain(result.vectors.iter().map(|e| e.entity.id.raw()))
.chain(result.documents.iter().map(|e| e.entity.id.raw()))
.chain(result.key_values.iter().map(|e| e.entity.id.raw()))
.collect();
result
.connections
.retain(|c| visible_ids.contains(&c.from_id) && visible_ids.contains(&c.to_id));
result.summary.total_entities = result.tables.len()
+ result.graph.nodes.len()
+ result.graph.edges.len()
+ result.vectors.len()
+ result.documents.len()
+ result.key_values.len();
}
fn require_visible<'a>(
scope: &'a dyn ReadFrame,
op: &'static str,
) -> RedDBResult<&'a HashSet<String>> {
match scope.visible_collections() {
Some(set) => Ok(set),
None => {
warn!(
target: "authorized_search",
op = op,
"refused: no visible-collections scope on frame"
);
Err(RedDBError::Query(format!(
"{op} requires an authenticated scope with visible_collections; \
none was attached to the runtime frame"
)))
}
}
}
fn constrain_collections(
requested: Option<Vec<String>>,
visible: &HashSet<String>,
) -> Option<Vec<String>> {
match requested {
None => {
let mut bounded: Vec<_> = visible.iter().cloned().collect();
bounded.sort();
Some(bounded)
}
Some(list) => {
let filtered: Vec<String> = list.into_iter().filter(|c| visible.contains(c)).collect();
Some(filtered)
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::api::RedDBOptions;
use crate::auth::Role;
use crate::runtime::statement_frame::test_support::FakeReadFrame;
use crate::runtime::RedDBRuntime;
fn rt() -> RedDBRuntime {
RedDBRuntime::with_options(RedDBOptions::in_memory()).expect("rt")
}
fn set(items: &[&str]) -> HashSet<String> {
items.iter().map(|s| s.to_string()).collect()
}
#[test]
fn execute_similar_refuses_without_scope() {
let rt = rt();
let frame = FakeReadFrame::without_scope();
let err = AuthorizedSearch::execute_similar(&rt, &frame, "orders", &[0.1], 1, 0.0)
.expect_err("refuses without scope");
assert!(format!("{err}").contains("requires an authenticated scope"));
}
#[test]
fn execute_similar_refuses_collection_outside_scope() {
let rt = rt();
let frame = FakeReadFrame::with_visible(set(&["orders"]));
let err = AuthorizedSearch::execute_similar(&rt, &frame, "secrets", &[0.1], 1, 0.0)
.expect_err("refuses out-of-scope collection");
assert!(format!("{err}").contains("not in the caller's visible scope"));
}
#[test]
fn execute_context_refuses_without_scope() {
let rt = rt();
let frame = FakeReadFrame::without_scope();
let err = AuthorizedSearch::execute_context(
&rt,
&frame,
SearchContextInput {
query: "x".into(),
field: None,
vector: None,
collections: None,
graph_depth: None,
graph_max_edges: None,
max_cross_refs: None,
follow_cross_refs: None,
expand_graph: None,
global_scan: None,
reindex: None,
limit: None,
min_score: None,
},
)
.expect_err("refuses without scope");
assert!(format!("{err}").contains("requires an authenticated scope"));
}
#[test]
fn constrain_collections_drops_out_of_scope_items() {
let visible = set(&["a", "b"]);
let got = constrain_collections(Some(vec!["a".into(), "c".into()]), &visible);
assert_eq!(got, Some(vec!["a".into()]));
assert_eq!(
constrain_collections(None, &visible),
Some(vec!["a".into(), "b".into()])
);
let _ = Role::Read;
}
use crate::runtime::{
ContextConnection, ContextConnectionType, ContextEntity, ContextGraphResult,
ContextSearchResult, ContextSummary, DiscoveryMethod,
};
use crate::storage::unified::entity::{EntityData, EntityKind, RowData, UnifiedEntity};
use crate::storage::unified::EntityId;
use proptest::prelude::*;
fn fake_entity(id: u64, collection: &str) -> UnifiedEntity {
UnifiedEntity::new(
EntityId::new(id),
EntityKind::TableRow {
table: std::sync::Arc::from(collection),
row_id: id,
},
EntityData::Row(RowData::new(Vec::new())),
)
}
fn fake_ctx_entity(id: u64, collection: &str) -> ContextEntity {
ContextEntity {
entity: fake_entity(id, collection),
score: 0.5,
discovery: DiscoveryMethod::GlobalScan,
collection: collection.to_string(),
}
}
fn empty_summary() -> ContextSummary {
ContextSummary {
total_entities: 0,
direct_matches: 0,
expanded_via_graph: 0,
expanded_via_cross_refs: 0,
expanded_via_vector_query: 0,
collections_searched: 0,
execution_time_us: 0,
tiers_used: Vec::new(),
entities_reindexed: 0,
}
}
fn build_result(rows: &[(u64, &str)]) -> ContextSearchResult {
let entities: Vec<ContextEntity> =
rows.iter().map(|(id, c)| fake_ctx_entity(*id, c)).collect();
ContextSearchResult {
query: "x".into(),
tables: entities.clone(),
graph: ContextGraphResult {
nodes: entities.clone(),
edges: Vec::new(),
},
vectors: entities.clone(),
documents: Vec::new(),
key_values: Vec::new(),
connections: vec![ContextConnection {
from_id: rows.first().map(|(id, _)| *id).unwrap_or(0),
to_id: rows.last().map(|(id, _)| *id).unwrap_or(0),
connection_type: ContextConnectionType::CrossRef("x".into()),
weight: 1.0,
}],
summary: empty_summary(),
}
}
proptest! {
#![proptest_config(ProptestConfig::with_cases(256))]
#[test]
fn every_result_row_is_in_visible_set(
row_collections in proptest::collection::vec("[a-z]{1,4}", 0..10),
visible in proptest::collection::hash_set("[a-z]{1,4}", 0..6),
) {
let rows: Vec<(u64, &str)> = row_collections
.iter()
.enumerate()
.map(|(i, c)| (i as u64 + 1, c.as_str()))
.collect();
let mut result = build_result(&rows);
post_filter_context_result(&mut result, &visible);
for e in result.tables.iter()
.chain(result.graph.nodes.iter())
.chain(result.graph.edges.iter())
.chain(result.vectors.iter())
.chain(result.documents.iter())
.chain(result.key_values.iter())
{
prop_assert!(visible.contains(&e.collection),
"leaked row collection={} not in visible={:?}",
e.collection, visible);
}
let visible_ids: HashSet<u64> = std::iter::empty()
.chain(result.tables.iter().map(|e| e.entity.id.raw()))
.chain(result.graph.nodes.iter().map(|e| e.entity.id.raw()))
.chain(result.graph.edges.iter().map(|e| e.entity.id.raw()))
.chain(result.vectors.iter().map(|e| e.entity.id.raw()))
.chain(result.documents.iter().map(|e| e.entity.id.raw()))
.chain(result.key_values.iter().map(|e| e.entity.id.raw()))
.collect();
for c in &result.connections {
prop_assert!(visible_ids.contains(&c.from_id) && visible_ids.contains(&c.to_id),
"dangling connection {} -> {} survived filter",
c.from_id, c.to_id);
}
}
}
#[test]
fn tenant_a_cannot_see_tenant_b_collection() {
let rt = rt();
let frame_a = FakeReadFrame::with_visible(set(&["a_orders"]));
let err = AuthorizedSearch::execute_similar(&rt, &frame_a, "b_orders", &[0.1], 1, 0.0)
.expect_err("tenant-A scope must refuse tenant-B collection");
assert!(format!("{err}").contains("not in the caller's visible scope"));
}
#[test]
fn search_context_refuses_all_out_of_scope_collections() {
let rt = rt();
let frame = FakeReadFrame::with_visible(set(&["a_orders"]));
let err = AuthorizedSearch::execute_context(
&rt,
&frame,
SearchContextInput {
query: "x".into(),
field: None,
vector: None,
collections: Some(vec!["b_orders".into(), "b_customers".into()]),
graph_depth: None,
graph_max_edges: None,
max_cross_refs: None,
follow_cross_refs: None,
expand_graph: None,
global_scan: None,
reindex: None,
limit: None,
min_score: None,
},
)
.expect_err("all-out-of-scope SEARCH CONTEXT must refuse");
let msg = format!("{err}");
assert!(
msg.contains("no requested collection is in the caller's visible scope"),
"expected scope-refusal, got: {msg}"
);
}
#[test]
fn auth_cache_stats_are_exposed_via_authstore() {
use crate::auth::store::AuthStore;
use crate::auth::AuthConfig;
let store = AuthStore::new(AuthConfig::default());
let stats0 = store.auth_cache_stats();
assert_eq!(stats0.hits + stats0.misses, 0);
let _ = store.visible_collections_for_scope(
None,
Role::Read,
"alice",
&vec!["orders".to_string()],
);
let stats1 = store.auth_cache_stats();
assert!(
stats1.misses >= 1,
"first lookup must record a miss, got {stats1:?}"
);
let _ = store.visible_collections_for_scope(
None,
Role::Read,
"alice",
&vec!["orders".to_string()],
);
let stats2 = store.auth_cache_stats();
assert!(
stats2.hits >= 1,
"second lookup must record a hit, got {stats2:?}"
);
let _ = stats2.hit_rate();
}
#[test]
fn visible_collections_cache_keeps_principals_separate() {
use crate::auth::privileges::{Action, GrantPrincipal, Resource};
use crate::auth::store::AuthStore;
use crate::auth::{AuthConfig, UserId};
let store = AuthStore::new(AuthConfig::default());
store.create_user("admin", "p", Role::Admin).unwrap();
store.create_user("alice", "p", Role::Read).unwrap();
store.create_user("bob", "p", Role::Read).unwrap();
store
.grant(
&UserId::platform("admin"),
Role::Admin,
GrantPrincipal::User(UserId::platform("alice")),
Resource::table_from_name("orders"),
vec![Action::Select],
false,
None,
)
.expect("grant alice orders");
let collections = vec!["orders".to_string()];
let alice = store.visible_collections_for_scope(None, Role::Read, "alice", &collections);
let bob = store.visible_collections_for_scope(None, Role::Read, "bob", &collections);
assert!(alice.contains("orders"));
assert!(
!bob.contains("orders"),
"bob must not reuse alice's visible-collections cache entry"
);
}
}
#[cfg(test)]
mod frame_regression_tests {
use super::*;
use crate::api::RedDBOptions;
use crate::auth::Role;
use crate::runtime::execution_context::{
snapshot_bundle, with_snapshot_bundle, SnapshotBundle,
};
use crate::runtime::statement_frame::test_support::FakeReadFrame;
use crate::storage::unified::entity::EntityData;
use reddb_types::Value;
fn text(runtime: &RedDBRuntime, frame: &dyn ReadFrame, limit: usize) -> Vec<String> {
let result = AuthorizedSearch::execute_text(
runtime,
frame,
"needle".into(),
None,
None,
None,
None,
Some(limit),
false,
)
.expect("text search");
result
.matches
.into_iter()
.map(|hit| {
let EntityData::Row(row) = hit.entity.data else {
panic!("row")
};
let Some(Value::Text(value)) = row.get_field("body") else {
panic!("text")
};
value.to_string()
})
.collect()
}
#[test]
fn search_text_uses_frame_allowlist_and_identity_before_limit() {
let rt = RedDBRuntime::with_options(RedDBOptions::in_memory()).expect("runtime");
for sql in [
"CREATE TABLE docs (owner TEXT, body TEXT)",
"CREATE TABLE secrets (owner TEXT, body TEXT)",
"INSERT INTO docs (owner,body) VALUES ('bob','needle denied'),('alice','needle visible')",
"INSERT INTO secrets (owner,body) VALUES ('alice','needle secret')",
"CREATE POLICY own ON docs FOR SELECT TO read USING (owner=CURRENT_USER())",
"ALTER TABLE docs ENABLE ROW LEVEL SECURITY",
] { rt.execute_query(sql).expect(sql); }
let mut frame = FakeReadFrame::with_visible(["docs".into()].into_iter().collect());
frame.snapshot = rt.current_snapshot();
let outer = SnapshotBundle {
auth: Some(("bob".into(), Role::Admin)),
tenant: Some("other".into()),
snapshot: None,
};
with_snapshot_bundle(&outer, || {
assert_eq!(text(&rt, &frame, 1), ["needle visible"]);
assert_eq!(text(&rt, &frame, 10), ["needle visible"]);
let restored = snapshot_bundle();
assert_eq!(restored.auth, outer.auth);
assert_eq!(restored.tenant, outer.tenant);
assert!(restored.snapshot.is_none());
});
frame.visible = Some(HashSet::new());
assert!(AuthorizedSearch::execute_text(
&rt,
&frame,
"needle".into(),
None,
None,
None,
None,
Some(10),
false
)
.is_err());
}
#[test]
fn search_text_reads_frame_snapshot_even_after_committed_update() {
let rt = RedDBRuntime::with_options(RedDBOptions::in_memory()).expect("runtime");
rt.execute_query("CREATE TABLE docs (id INT PRIMARY KEY, body TEXT)")
.expect("table");
rt.execute_query("INSERT INTO docs (id,body) VALUES (1,'needle old')")
.expect("row");
let mut frame = FakeReadFrame::with_visible(["docs".into()].into_iter().collect());
frame.snapshot = rt.current_snapshot();
rt.execute_query("UPDATE docs SET body='needle new' WHERE id=1")
.expect("update");
assert_eq!(text(&rt, &frame, 10), ["needle old"]);
frame.snapshot = rt.current_snapshot();
assert_eq!(text(&rt, &frame, 10), ["needle new"]);
}
#[test]
fn search_similar_filters_policy_before_topk_and_uses_frame_identity() {
let rt = RedDBRuntime::with_options(RedDBOptions::in_memory()).expect("runtime");
for sql in [
"CREATE VECTOR embeddings DIM 2 METRIC cosine",
"INSERT INTO embeddings VECTOR (dense,content) VALUES ([1.0,0.0],'denied') WITH METADATA (owner='bob')",
"INSERT INTO embeddings VECTOR (dense,content) VALUES ([0.8,0.6],'visible') WITH METADATA (owner='alice')",
"CREATE POLICY own ON VECTORS OF embeddings FOR SELECT TO read USING (metadata.owner=CURRENT_USER())",
"ALTER TABLE embeddings ENABLE ROW LEVEL SECURITY",
] { rt.execute_query(sql).expect(sql); }
let mut frame = FakeReadFrame::with_visible(["embeddings".into()].into_iter().collect());
frame.snapshot = rt.current_snapshot();
with_snapshot_bundle(
&SnapshotBundle {
auth: Some(("bob".into(), Role::Admin)),
..Default::default()
},
|| {
let hits = AuthorizedSearch::execute_similar(
&rt,
&frame,
"embeddings",
&[1.0, 0.0],
1,
0.0,
)
.expect("similar");
assert_eq!(hits.len(), 1);
let EntityData::Vector(data) = &hits[0].entity.data else {
panic!("vector")
};
assert_eq!(data.content.as_deref(), Some("visible"));
},
);
}
#[test]
fn search_scope_preserves_own_savepoint_writes_and_aborts() {
use crate::runtime::statement_frame::{StatementExecutionFrame, StatementIdentity};
let rt = RedDBRuntime::with_options(RedDBOptions::in_memory()).expect("runtime");
for sql in [
"CREATE TABLE docs (id INT, body TEXT)",
"BEGIN",
"INSERT INTO docs (id,body) VALUES (1,'needle parent')",
"SAVEPOINT child",
"INSERT INTO docs (id,body) VALUES (2,'needle child')",
] {
rt.execute_query(sql).expect(sql);
}
let scope = {
let frame = StatementExecutionFrame::build(
&rt,
StatementIdentity::Text("SEARCH TEXT 'needle' IN docs"),
)
.expect("frame");
let _installed = frame.install(&rt);
let mut scope = rt.ai_scope();
scope.visible_collections = Some(["docs".into()].into_iter().collect());
scope
};
let mut rows = text(&rt, &scope, 10);
rows.sort();
assert_eq!(rows, ["needle child", "needle parent"]);
rt.execute_query("ROLLBACK TO SAVEPOINT child")
.expect("rollback child");
assert_eq!(text(&rt, &scope, 10), ["needle parent"]);
rt.execute_query("ROLLBACK").expect("rollback parent");
assert!(text(&rt, &scope, 10).is_empty());
}
}