diff --git a/aimdb-core/CHANGELOG.md b/aimdb-core/CHANGELOG.md index 0938bf41..981610de 100644 --- a/aimdb-core/CHANGELOG.md +++ b/aimdb-core/CHANGELOG.md @@ -7,6 +7,26 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Changed (breaking, API) + +- **`ConnectorConfig` gains `record_index: Option`**, the id of the + record an outbound publish comes from — its registration index, the same + `record_id` `AimDb::list_records` reports. A topic alone cannot identify the + record, since several records may publish on one topic (or on topics a + `TopicProvider` picks per value), so a connector enforcing per-record + authorization needs the index at publish time; the WebSocket connector uses it + to gate delivery against each client's granted records. The struct is not + `#[non_exhaustive]`, so code building a `ConnectorConfig` with a struct + literal must add the field (or use `..Default::default()`); it defaults to + `None`. + +- **`AimDb::collect_outbound_routes` stamps each route with its record index**, + appending `("record_index", "")` to `OutboundRoute::config`, and + `ConnectorConfig::from_query` parses that pair into `record_index` rather than + passing it on in `protocol_options`. `record_index` is therefore a reserved + config key: a value set through `with_config("record_index", …)` is + overridden by the stamped one, which is appended last. + ## [2.0.0] - 2026-09-18 ### Added diff --git a/aimdb-core/src/builder.rs b/aimdb-core/src/builder.rs index 635e8773..80710d55 100644 --- a/aimdb-core/src/builder.rs +++ b/aimdb-core/src/builder.rs @@ -64,6 +64,9 @@ struct RecordEntry { pub struct AimDbInner { /// Record entries, indexed by `RecordId`. Order matches registration /// order. Immutable after build(). + /// + /// Per-client authorization is based on index of RecordEntry in `storage` + /// Never reorder, remove, or `swap_remove` entries as these could break current logic. storages: Vec, /// Name → RecordId lookup (control plane) @@ -1222,7 +1225,14 @@ impl AimDb { pub fn collect_outbound_routes(&self, scheme: &str) -> Vec { let mut routes = Vec::new(); - for entry in &self.inner.storages { + for (i, entry) in self.inner.storages.iter().enumerate() { + // i and RecordId must match + debug_assert_eq!( + self.inner.by_key.get(&entry.key).map(|id| id.index()), + Some(i), + "record storage order diverges from RecordId for key {}", + entry.key.as_str() + ); let outbound_links = entry.record.outbound_connectors(); for link in outbound_links { @@ -1231,11 +1241,15 @@ impl AimDb { continue; } + // config must carry the record index + let mut config = link.config.clone(); + config.push(("record_index".to_string(), i.to_string())); + // Create the fused source using the stored factory routes.push(OutboundRoute { topic: link.url.resource_id().to_string(), source: link.create_source(self), - config: link.config.clone(), + config, }); } } diff --git a/aimdb-core/src/transport.rs b/aimdb-core/src/transport.rs index 94aa31e2..da5cf6a5 100644 --- a/aimdb-core/src/transport.rs +++ b/aimdb-core/src/transport.rs @@ -33,6 +33,11 @@ pub struct ConnectorConfig { /// Protocol-specific options as key-value pairs /// Allows custom configuration without polluting the base struct pub protocol_options: Vec<(String, String)>, + + /// The index of the record key that is setup together with the outbound route. + /// The record key is unique, and its order in `AimDb::inner.storage` is immutable + /// so this `record_index` could be used for `O(1)` lookup. + pub record_index: Option, } impl Default for ConnectorConfig { @@ -40,6 +45,7 @@ impl Default for ConnectorConfig { Self { timeout_ms: Some(5000), protocol_options: Vec::new(), + record_index: None, } } } @@ -51,8 +57,9 @@ impl ConnectorConfig { /// per-route configuration through to [`Connector::publish`] without changing /// the `publish` signature. /// - /// Only the protocol-agnostic `timeout_ms` is lifted into the typed field; - /// every other key is passed through verbatim in + /// Only the protocol-agnostic `timeout_ms` and `record_index` (stamped by + /// `AimDb::collect_outbound_routes`; the last occurrence wins) are lifted + /// into typed fields; every other key is passed through verbatim in /// [`protocol_options`](ConnectorConfig::protocol_options) for the /// connector to interpret with its own defaults. pub fn from_query(query: &[(String, String)]) -> ConnectorConfig { @@ -64,6 +71,11 @@ impl ConnectorConfig { cfg.timeout_ms = Some(n); } } + "record_index" => { + if let Ok(i) = v.parse::() { + cfg.record_index = Some(i); + } + } _ => cfg.protocol_options.push((k.clone(), v.clone())), } } diff --git a/aimdb-websocket-connector/CHANGELOG.md b/aimdb-websocket-connector/CHANGELOG.md index 1aedaeee..ab79aa10 100644 --- a/aimdb-websocket-connector/CHANGELOG.md +++ b/aimdb-websocket-connector/CHANGELOG.md @@ -7,6 +7,118 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Changed (breaking, API) + +- **Read grants are record-key patterns, not topic patterns.** + `Permissions::subscribe_patterns` is renamed to `read_patterns` and + `Permissions::can_subscribe(topic)` is renamed `can_read(key)`. The rename + carries a sublte change: the read patterns are now matched against **record + keys** — the string passed to `configure::(key, ...)` — rather than against + the ws topic registered at `link_to("ws://...")`. `write_patterns` and `can_write` are + unchanged and remain **topic**-based. An operator migrating a config + replaces each topic grant with the key(s) of the records publishing there; + where the two namespaces coincide (`link_to("ws://cfg")` on record `cfg`) no + change is needed. + +- **`AuthHandler::authorize_subscribe`, `authorize_query` and `authorize_list` + are removed.** Authorization is resolved once, at the HTTP upgrade, into a + per-client record bitset; the three read paths consult that bitset instead of + calling back into the handler. A handler that overrode any of them must move + the logic into `authenticate`, which now returns the grants that become the + bitset. `authorize_write` is unaffected. + + **Read grants are now fixed for the lifetime of a connection.** Previously a + handler overriding these hooks could consult external or mutable state (an ACL + service, a revocation list) on every `sub`, `record.query` and `record.list`, + and so refuse *new* read requests on an open connection once a grant changed. + That per-request re-check is gone: nothing on the read path calls back into + the handler after `authenticate`. With the default handler nothing changes — + grants were already fixed at the upgrade. Live subscriptions were never + re-checked per event, before or after this change, so revoking a client's + reads still means closing its connection and letting it re-authenticate. + +- **`AuthHandler::authorize_query_record` added**, a `record.query` post-filter + receiving the handler's rows and returning the subset the client may read. + It defaults to `Permissions::can_read` on each row's key, so implementors need + it only to narrow further. This is deliberately pattern-based rather than + bitset-based: history outlives configuration, so a query may legitimately name + a record the running server no longer registers. + +- **`ClientInfo` carries `record_perms: Arc`**, the resolved bitset, + alongside the `Permissions` it was built from. `RecordsBits` is public + (`new`, `set`, `is_allowed`, `len`, `is_empty`, + `resolve_permissions`); one bit per registered record, indexed by record id — + registration order in the builder, which is also the index + `AimDb::list_records` reports. + +- **`ClientManager::subscribe` and `broadcast` signatures changed.** + `subscribe(pattern)` becomes `subscribe(pattern, Arc)` — a + subscription now carries the grants it delivers under — and + `broadcast(topic, payload)` becomes `broadcast(topic, record_index, payload)`, + since the topic alone no longer identifies which record a message came from. + +- **`SnapshotProvider::snapshots` returns `Vec<(usize, String, Vec)>`** + instead of `Vec<(String, Vec)>`; the added `usize` is the record id the + cached value belongs to. Two records publishing on one topic now yield two + entries rather than one overwriting the other. + +- **`ConnectorConfig` carry an additional attr `record_index: Option` + to `WsBusSink`**, so record index could join topic at outbound routes. + +### Changed (breaking, wire) + +- **A late-join burst may carry several `snap` frames for the same topic** — one + per granted record publishing there — where previously a topic produced at + most one. Each rides its own `seq`, and `last` still closes the burst, so a + client tracking sequence numbers is unaffected; a client indexing snapshots by + topic must expect collisions. + +- **`record.query`'s `total` is the number of rows returned**, after + authorization, rather than the handler's own match count. The trait doc for + `QueryHandler::handle_query` is updated to say so: a client is never told how + many rows it was not allowed to see. The built-in persistence handler already + returned `records.len()`, so only custom handlers change behaviour. + + However, there is a limit to this design: in case `total` < `limit`, it does not mean + that there is no more record, but records returned from `QueryHandlerFn` are filtered by grants. + A better design is to resolve search name against grants then dedicate the filtering + to `QueryHandlerFn` instead of let it stay in `WsSession`. + +- **`record.query` distinguishes "not permitted" from "nothing matched".** + `denied` is returned only when the client holds no read grants at all; + a client with grants that match no stored rows gets `{"records": [], "total": 0}`. + Querying a record the server no longer registers is allowed if the grant + covers its key, so history survives a record's retirement from the config. + +- **Subscribing to a topic no records publish on now succeeds and stays silent** + rather than being refused. Topics are not part of the authorization model any + more, so the server cannot tell an unregistered topic from one whose records + the client may not read. Only a client holding no read grants at all is + refused with `denied`; a client whose grants match no registered record + subscribes and stays silent, as does any client on a server with no records. + +- **`record.list` returns only the records the client may read.** Previously the + full database was enumerable by any authenticated client. + +### Fixed + +- **Late-join snapshots no longer lose a record when two share a topic.** The + snapshot cache was keyed by topic, so the most recent publisher overwrote the + other's cached value and a late-joining client received only one of them — + and, once grants became per record, sometimes neither. It is now keyed by + `(record id, topic)`. + +### Security + +- **A read grant no longer leaks records it does not name.** Grants lived in + topic space while records are keyed independently, so a client granted one + record's topic received every record publishing there — including records it + held no grant for, over `event` frames, late-join snapshots, `record.list` and + `record.query` alike. Records whose topic comes from a `TopicProvider` made + this unbounded, since the topic is chosen per value at runtime. Grants are now + resolved against record keys at the upgrade and enforced at every delivery + point, so a record's data reaches only clients granted that record. + ## [0.3.0] - 2026-09-18 ### Added diff --git a/aimdb-websocket-connector/src/lib.rs b/aimdb-websocket-connector/src/lib.rs index 1ad31abd..5710c4bb 100644 --- a/aimdb-websocket-connector/src/lib.rs +++ b/aimdb-websocket-connector/src/lib.rs @@ -125,7 +125,7 @@ pub type WebSocketConnector = server::builder::WebSocketConnectorBuilder; #[cfg(feature = "server")] pub use server::auth::{ - AuthError, AuthHandler, AuthRequest, ClientId, ClientInfo, NoAuth, Permissions, + AuthError, AuthHandler, AuthRequest, ClientId, ClientInfo, NoAuth, Permissions, RecordsBits, }; #[cfg(feature = "server")] pub use server::client_manager::ClientManager; diff --git a/aimdb-websocket-connector/src/server/auth.rs b/aimdb-websocket-connector/src/server/auth.rs index 297b6ae5..c7f2f0aa 100644 --- a/aimdb-websocket-connector/src/server/auth.rs +++ b/aimdb-websocket-connector/src/server/auth.rs @@ -2,16 +2,11 @@ //! //! The [`AuthHandler`] trait provides pluggable auth hooks for: //! -//! 1. **Connection upgrade** — `authenticate()`: decide whether to accept the WebSocket -//! handshake and assign per-client permissions. -//! 2. **Topic subscriptions** — `authorize_subscribe()`: gate which topics a client can -//! receive data from. -//! 3. **Inbound writes** — `authorize_write()`: gate which topics a client may write to. -//! 4. **Historical reads** — `authorize_query()`: gate the `record.query` pattern. -//! 5. **Introspection** — `authorize_list()`: gate which `record.list` rows a client sees. -//! -//! (4) and (5) default to (2), so overriding `authorize_subscribe` governs all -//! three read paths. +//! 1. **Connection upgrade** — `authenticate()`: resolve per-client permissions into bitmasks +//! which checked for before message broadcasting. +//! `authenticate()` does not gate topic subscription, so clients could claim unregistered topics, +//! and receive nothing during their lifetime. +//! 2. **Inbound writes** — `authorize_write()`: gate which topic a client may write to. //! //! The default implementation ([`NoAuth`]) allows all operations. @@ -22,6 +17,7 @@ use std::sync::Arc; use core::future::Future; use core::pin::Pin; +use aimdb_core::remote::QueryRecord; use axum::http::HeaderMap; // ════════════════════════════════════════════════════════════════════ @@ -44,53 +40,58 @@ pub struct ClientInfo { pub id: ClientId, pub remote_addr: SocketAddr, pub permissions: Permissions, + /// Bitsets represent accessibility of record at index `i` + /// Records maintain the same order as in `AimdDb.inner` + pub record_perms: Arc, } /// Per-client permission set assigned during authentication. /// -/// Each field is a list of topic *patterns* (supporting `*` and `#` wildcards +/// Each field is a list of record key *patterns* (supporting `*` and `#` wildcards /// as defined by [`aimdb_core::topic_matches`]). /// /// An empty `Vec` means *"no access"*. Use `["#"]` for unrestricted access. #[derive(Debug, Clone, Default)] pub struct Permissions { - /// Topic patterns the client may subscribe to. - pub subscribe_patterns: Vec, + /// Record name patterns the client may read from. + pub read_patterns: Vec, /// Topic patterns the client may write to. pub write_patterns: Vec, } impl Permissions { - /// Creates a permission set that grants full access to everything. + /// Creates a permission set that grants full access to all records. pub fn allow_all() -> Self { Self { - subscribe_patterns: vec!["#".to_string()], + read_patterns: vec!["#".to_string()], write_patterns: vec!["#".to_string()], } } - /// Returns `true` if the client is allowed to subscribe to `topic`. + /// Returns `true` if the client is allowed to access to record with `key`. /// - /// `topic` is the client's requested subscription, which may itself be a - /// wildcard — so this asks **pattern containment** + /// `key` is a registered record key or a pattern of record key — + /// so this asks **pattern containment** + /// Currently, the server passes concrete keys to setup permissions bitmasks + /// for a client during http upgrade. /// ([`pattern_contains`](aimdb_core::pattern_contains)): /// does a granted pattern cover the *whole* requested pattern? Plain /// [`topic_matches`](aimdb_core::topic_matches) would let a one-level grant /// (`sensors.*`) admit an all-levels request (`sensors.#`) by having the /// `*` swallow the `#`, silently widening the grant. For a concrete request - /// `pattern_contains` collapses to `topic_matches`, so exact subscribes are + /// `pattern_contains` collapses to `topic_matches`, so exact permissions are /// unaffected. - pub fn can_subscribe(&self, topic: &str) -> bool { - self.subscribe_patterns + pub fn can_read(&self, key: &str) -> bool { + self.read_patterns .iter() - .any(|p| aimdb_core::pattern_contains(p, topic)) + .any(|p| aimdb_core::pattern_contains(p, key)) } - /// Returns `true` if the client is allowed to write to `topic`. + /// Returns `true` if the client is allowed to write to a record given with key. /// - /// Writes target a single concrete record, so `topic` is never a wildcard + /// Writes target a single concrete record, so `key` is never a wildcard /// here and plain [`topic_matches`](aimdb_core::topic_matches) is the right - /// check (a wildcard write topic would resolve to no record downstream). + /// check (a wildcard write key would resolve to no record downstream). pub fn can_write(&self, topic: &str) -> bool { self.write_patterns .iter() @@ -98,6 +99,85 @@ impl Permissions { } } +/// Bit mask to store records that client has access to. +/// +/// Index `i` of the mask represent access to record id `i`. +/// The record having `id` assigned incrementally, according to [`aimdb_core::builder::AimDbInner`] +/// As the crate does not use any dependency for bit set, we craft one +/// from vector of `u8` for finer granularity. +/// +/// `i` is proportionate to bit significance, so a block mask could be easily constructed +/// using `1u8 << i` +#[derive(Debug, Clone)] +pub struct RecordsBits { + length: usize, + masks: Vec, +} + +impl RecordsBits { + pub fn new(length: usize) -> Self { + let blocks = length.div_ceil(8); + let masks: Vec = vec![0u8; blocks]; + Self { length, masks } + } + + pub fn len(&self) -> usize { + self.length + } + + /// Simply check for length of the bitmap + pub fn is_empty(&self) -> bool { + self.length == 0 + } + + /// Create from Permissions + pub fn resolve_permissions(records: &[String], permissions: &Permissions) -> Self { + let mut records_bits = Self::new(records.len()); + records.iter().enumerate().for_each(|(i, key)| { + if permissions.can_read(key.as_str()) { + let _ = records_bits.set(i); + }; + }); + records_bits + } + + /// Checks whether record at `index` is accessible. + /// Out-of-index index return `false`, so no panic. + pub fn is_allowed(&self, index: usize) -> bool { + if index >= self.length { + return false; + }; + + // Bit index inside the block, + let block_index = Self::block_index(index); + let offset = Self::offset(index); + let mask = 1u8 << offset; + self.masks[block_index] & mask != 0 + } + + /// Set record at `index` accessible, out-of-index returns `false` + pub fn set(&mut self, index: usize) -> bool { + if index >= self.length { + return false; + }; + + // Bit index inside the block, + let block_index = Self::block_index(index); + let offset = Self::offset(index); + let mask = 1u8 << offset; + self.masks[block_index] |= mask; + true + } + + fn block_index(index: usize) -> usize { + index / 8 + } + + fn offset(index: usize) -> usize { + index % 8 + } +} + /// Context provided to [`AuthHandler::authenticate`] during WebSocket upgrade. #[derive(Debug)] pub struct AuthRequest { @@ -168,17 +248,6 @@ pub trait AuthHandler: Send + Sync + 'static { request: &'a AuthRequest, ) -> Pin> + Send + 'a>>; - /// Called before allowing a topic subscription. - /// - /// The default implementation delegates to [`Permissions::can_subscribe`]. - fn authorize_subscribe<'a>( - &'a self, - client: &'a ClientInfo, - topic: &'a str, - ) -> Pin + Send + 'a>> { - Box::pin(async move { client.permissions.can_subscribe(topic) }) - } - /// Called before routing an inbound write to a producer. /// /// The default implementation delegates to [`Permissions::can_write`]. @@ -190,35 +259,17 @@ pub trait AuthHandler: Send + Sync + 'static { Box::pin(async move { client.permissions.can_write(topic) }) } - /// Called before serving a `record.query` (historical read). - /// - /// `pattern` is the query's `name`, possibly wildcarded — so this is the - /// containment check [`authorize_subscribe`](Self::authorize_subscribe) - /// performs, which it delegates to by default. A query omitting `name` asks - /// for `"*"`, which a narrower grant does not contain: it fails closed. - fn authorize_query<'a>( + /// Keeps records having key matched grants + /// `records` are from a [QueryHandlerFn](aimdb_core::remote::QueryHandlerFn) + fn authorize_query_record<'a>( &'a self, client: &'a ClientInfo, - pattern: &'a str, - ) -> Pin + Send + 'a>> { - self.authorize_subscribe(client, pattern) - } - - /// Called for each `record.list` row; `false` drops the row and the call - /// still succeeds. Defaults to - /// [`authorize_subscribe`](Self::authorize_subscribe). - /// - /// `record_key` is the row's database key, *not* its WebSocket topic. The - /// two coincide under `link_to("ws://")`, but a `TopicProvider` that - /// computes the topic per value has no single topic to check against — - /// grant the record key itself (`["sensors.#", "inject"]`) to keep such a - /// record introspectable. - fn authorize_list<'a>( - &'a self, - client: &'a ClientInfo, - record_key: &'a str, - ) -> Pin + Send + 'a>> { - self.authorize_subscribe(client, record_key) + records: Vec, + ) -> Vec { + records + .into_iter() + .filter(|r| client.permissions.can_read(&r.topic)) + .collect() } } @@ -243,23 +294,62 @@ pub(crate) type DynAuthHandler = Arc; #[cfg(test)] mod tests { - use super::Permissions; + use super::{Permissions, RecordsBits}; fn perms(subscribe: &[&str]) -> Permissions { Permissions { - subscribe_patterns: subscribe.iter().map(|s| s.to_string()).collect(), + read_patterns: subscribe.iter().map(|s| s.to_string()).collect(), write_patterns: Vec::new(), } } + #[test] + fn records_bits_works() { + let mut records_bits = RecordsBits::new(20); + + assert_eq!(records_bits.masks.len(), 3); + + records_bits.set(11); + assert!(records_bits.is_allowed(11)); + assert!(!records_bits.is_allowed(10)); + + // Out of index + assert!(!records_bits.set(20)); + } + + #[test] + fn records_bits_boundaries() { + // Length on a byte boundary: exactly two blocks, no spare one + let mut records_bits = RecordsBits::new(16); + assert_eq!(records_bits.masks.len(), 2); + + // First and last bit of each block + for i in [0, 7, 8, 15] { + assert!(records_bits.set(i)); + assert!(records_bits.is_allowed(i)); + } + assert!(!records_bits.is_allowed(1)); + assert!(!records_bits.is_allowed(14)); + + // `len` is the first out-of-range index, for both set and read + assert!(!records_bits.set(16)); + assert!(!records_bits.is_allowed(16)); + + // Zero records: nothing to set or read, no panic + let mut empty = RecordsBits::new(0); + assert!(empty.masks.is_empty()); + assert!(!empty.set(0)); + assert!(!empty.is_allowed(0)); + } + #[test] fn subscribe_grant_is_not_widened_by_a_wildcard_request() { // A one-level grant must not admit an all-levels request (the `*`-eats-`#` // escalation): `sensors.*` covers one level, `sensors.#` covers all. let p = perms(&["sensors.*"]); - assert!(p.can_subscribe("sensors.temp")); // in scope - assert!(!p.can_subscribe("sensors.#")); // escalation — denied - assert!(!p.can_subscribe("sensors.temp.vienna")); // deeper — denied + assert!(p.can_read("sensors.temp")); // in scope + assert!(!p.can_read("sensors.#")); // escalation — denied + assert!(!p.can_read("sensors.temp.vienna")); // deeper — denied } #[test] @@ -268,21 +358,54 @@ mod tests { // of the subtree: the `#` used to short-circuit the whole match, so every // segment after it was ignored and this grant behaved like `tenant.#`. let p = perms(&["tenant.#.secret"]); - assert!(p.can_subscribe("tenant.secret")); - assert!(p.can_subscribe("tenant.a.b.secret")); - assert!(!p.can_subscribe("tenant.public")); - assert!(!p.can_subscribe("tenant.a.b.public")); - assert!(!p.can_subscribe("tenant.#")); // no escalation to the subtree + assert!(p.can_read("tenant.secret")); + assert!(p.can_read("tenant.a.b.secret")); + assert!(!p.can_read("tenant.public")); + assert!(!p.can_read("tenant.a.b.public")); + assert!(!p.can_read("tenant.#")); // no escalation to the subtree } #[test] fn subscribe_allows_requests_the_grant_actually_covers() { - assert!(perms(&["#"]).can_subscribe("sensors.#")); + assert!(perms(&["#"]).can_read("sensors.#")); let p = perms(&["sensors.#"]); - assert!(p.can_subscribe("sensors.#")); - assert!(p.can_subscribe("sensors.temp.#")); - assert!(p.can_subscribe("sensors.temp")); + assert!(p.can_read("sensors.#")); + assert!(p.can_read("sensors.temp.#")); + assert!(p.can_read("sensors.temp")); // Out of the granted subtree stays denied. - assert!(!p.can_subscribe("commands.#")); + assert!(!p.can_read("commands.#")); + } + + #[test] + fn permissions_bitmask_from_grant_and_perms() { + // Seed permissions + let perms = perms(&["home.#", "front.#.test", "garage.#"]); + let mut records: Vec = Vec::new(); + records.extend( + [ + "home.1.1", + "back.1", + "front.1", + "front.2.test", + "back.2", + "garage.inner.1", + "garage.outer.1", + "storage.1", + ] + .iter() + .map(|s| s.to_string()), + ); + + let record_perms = RecordsBits::resolve_permissions(&records, &perms); + + assert_eq!(record_perms.len(), records.len()); + assert!(record_perms.is_allowed(0)); + assert!(!record_perms.is_allowed(1)); + assert!(!record_perms.is_allowed(2)); + assert!(record_perms.is_allowed(3)); + assert!(!record_perms.is_allowed(4)); + assert!(record_perms.is_allowed(5)); + assert!(record_perms.is_allowed(6)); + assert!(!record_perms.is_allowed(7)); } } diff --git a/aimdb-websocket-connector/src/server/builder.rs b/aimdb-websocket-connector/src/server/builder.rs index c882cff7..d8149ee7 100644 --- a/aimdb-websocket-connector/src/server/builder.rs +++ b/aimdb-websocket-connector/src/server/builder.rs @@ -235,8 +235,8 @@ impl WebSocketConnectorBuilder { /// /// Without this, `record.query` delegates to the `QueryHandlerFn` that /// `aimdb-persistence::with_persistence` registers in Extensions; with - /// neither, clients get `not_found`. Either way the requested pattern passes - /// [`AuthHandler::authorize_query`] first. + /// neither, clients get `not_found`. Rows outside of client's permissions + /// and query pattern are not returned. pub fn with_query_handler(mut self, handler: impl QueryHandler + 'static) -> Self { self.query_handler = Some(Arc::new(handler)); self @@ -306,6 +306,15 @@ impl ConnectorBuilder for WebSocketConnectorBuilder { // the dispatch the name lookup to stamp onto core's rows. let schema_by_type = Arc::new(self.streamable_registry.schema_by_type_id()); + // List of record key from inner db, order maintained as registration order + // This list is immutable after this `build` completed + let records: Arc> = Arc::new( + db.list_records() + .iter() + .map(|m| m.record_key.clone()) + .collect(), + ); + // ── Shared dispatch (one Arc per server) ─── let dispatch: Arc = Arc::new(WsDispatch { db: db.clone(), @@ -339,6 +348,7 @@ impl ConnectorBuilder for WebSocketConnectorBuilder { max_clients: self.max_clients.max(1), max_subs_per_connection: self.max_subs_per_connection.max(1), started_at: Instant::now(), + records, }; let additional = self.additional_routes.clone(); let server_future = @@ -356,16 +366,17 @@ impl ConnectorBuilder for WebSocketConnectorBuilder { // Dynamic snapshot provider backed by the shared Mutex // ════════════════════════════════════════════════════════════════════ +/// Wrapper for SnapshotCache struct DynMapSnapshot(SnapshotCache); impl SnapshotProvider for DynMapSnapshot { - fn snapshots(&self, pattern: &str) -> Vec<(String, Vec)> { + fn snapshots(&self, pattern: &str) -> Vec<(usize, String, Vec)> { let Ok(map) = self.0.lock() else { return Vec::new(); }; map.iter() - .filter(|(topic, _)| topic_matches(pattern, topic)) - .map(|(topic, bytes)| (topic.clone(), bytes.clone())) + .filter(|((_, topic), _)| topic_matches(pattern, topic)) + .map(|((record_id, topic), bytes)| (*record_id, topic.clone(), bytes.clone())) .collect() } } diff --git a/aimdb-websocket-connector/src/server/client_manager.rs b/aimdb-websocket-connector/src/server/client_manager.rs index 830a059c..7c9c40f6 100644 --- a/aimdb-websocket-connector/src/server/client_manager.rs +++ b/aimdb-websocket-connector/src/server/client_manager.rs @@ -20,6 +20,8 @@ use aimdb_core::{topic_matches, BoxStream, Payload, SubUpdate}; use dashmap::DashMap; use tokio::sync::mpsc; +use crate::server::auth::RecordsBits; + use super::auth::ClientId; /// One live subscription: a wildcard pattern + the channel feeding its stream. @@ -34,6 +36,11 @@ struct SubEntry { /// this, a drop *here* (upstream of where the pump assigns `seq`) would be /// silent, and a slow fan-out consumer would under-report its loss. dropped: AtomicU64, + + /// A bit mask encodes what record that the client could read from. + /// `SubEntry` is created when a client subscribes to a topic, so it ultimately + /// could carry a bit of `ClientInfo` + record_perms: Arc, } /// Drop guard for SubEntry @@ -128,7 +135,11 @@ impl ClientManager { /// through a guard that the stream owns. /// The next matching [`broadcast`](Self::broadcast) keeps trying to lazily prune the entry. /// This could serve as a safety net to make sure nothing leaks. - pub fn subscribe(&self, pattern: &str) -> (u64, BoxStream<'static, SubUpdate>) { + pub fn subscribe( + &self, + pattern: &str, + record_perms: Arc, + ) -> (u64, BoxStream<'static, SubUpdate>) { let id = self.next_sub.fetch_add(1, Ordering::Relaxed); let (tx, rx) = mpsc::channel::(self.sub_capacity); self.subs.insert( @@ -137,6 +148,7 @@ impl ClientManager { pattern: pattern.to_string(), tx, dropped: AtomicU64::new(0), + record_perms, }, ); // A drop guard for RAII, thankfully self.subs is already Arc<_> @@ -160,7 +172,11 @@ impl ClientManager { /// A full channel drops the update (slow-client protection) but records it on /// the subscription's `dropped` counter, folded into the next delivered /// update's `skipped` so the loss still surfaces as a `seq` gap. - pub async fn broadcast(&self, topic: &str, payload_bytes: &[u8]) { + pub async fn broadcast(&self, topic: &str, record_index: usize, payload_bytes: &[u8]) { + // A record could comes with multiple topic, or none + // But a client may not have access to some records due to server's authorization process + // Broadcasting must check for allowed records of each client + let payload = Payload::from(payload_bytes); let tag: Arc = Arc::from(topic); let mut dead: Vec = Vec::new(); @@ -168,6 +184,12 @@ impl ClientManager { if !topic_matches(&entry.pattern, topic) { continue; } + + // Do not broadcast for non-authorized records + if !entry.record_perms.is_allowed(record_index) { + continue; + } + // Carry any drops accumulated since the last delivered update, so a // broadcast-stage loss rides this update's `skipped` into a `seq` // gap. Take them now; restore on failure so nothing is lost. @@ -221,12 +243,20 @@ mod tests { use super::*; use futures_util::StreamExt; + /// A simple record permissions that allow client to read record having index 0 + fn simple_record_perms() -> Arc { + let mut record_pems = RecordsBits::new(1); + let _ = record_pems.set(0); + Arc::new(record_pems) + } + #[tokio::test] async fn broadcast_reaches_matching_subscriptions() { let mgr = ClientManager::new(256); - let (_id, mut stream) = mgr.subscribe("sensors.#"); + let record_perms = simple_record_perms(); + let (_id, mut stream) = mgr.subscribe("sensors.#", record_perms); - mgr.broadcast("sensors.temp.vienna", b"22.5").await; + mgr.broadcast("sensors.temp.vienna", 0, b"22.5").await; // Delivery is the raw payload tagged with the real topic — even for the // wildcard sub; the envelope is the per-connection codec's job. @@ -239,8 +269,9 @@ mod tests { async fn non_matching_topic_is_not_delivered() { use futures_util::FutureExt; let mgr = ClientManager::new(256); - let (_id, mut stream) = mgr.subscribe("commands.#"); - mgr.broadcast("sensors.temp", b"22.5").await; + let record_perms = simple_record_perms(); + let (_id, mut stream) = mgr.subscribe("commands.#", record_perms); + mgr.broadcast("sensors.temp", 0, b"22.5").await; // Nothing queued: the next() future is not ready. assert!(stream.next().now_or_never().is_none()); } @@ -251,11 +282,12 @@ mod tests { // dropped, but the loss must ride the next delivered update's `skipped` // so it becomes a `seq` gap downstream (not a silent hole). let mgr = ClientManager::new(1); - let (_id, mut stream) = mgr.subscribe("#"); + let record_perms = simple_record_perms(); + let (_id, mut stream) = mgr.subscribe("#", record_perms); - mgr.broadcast("t", b"1").await; // fills the one slot - mgr.broadcast("t", b"2").await; // full → dropped (counter = 1) - mgr.broadcast("t", b"3").await; // full → dropped (counter = 2) + mgr.broadcast("t", 0, b"1").await; // fills the one slot + mgr.broadcast("t", 0, b"2").await; // full → dropped (counter = 1) + mgr.broadcast("t", 0, b"3").await; // full → dropped (counter = 2) // First delivery is the update that got through, lossless. let first = stream.next().await.expect("first update"); @@ -264,7 +296,7 @@ mod tests { // With the slot now free, the next broadcast is delivered and carries the // two drops that happened while the channel was full. - mgr.broadcast("t", b"4").await; + mgr.broadcast("t", 0, b"4").await; let second = stream.next().await.expect("second update"); assert_eq!(&second.data[..], b"4"); assert_eq!( @@ -276,8 +308,11 @@ mod tests { #[tokio::test] async fn fan_out_to_n_subscribers() { let mgr = ClientManager::new(256); - let mut streams: Vec<_> = (0..5).map(|_| mgr.subscribe("#").1).collect(); - mgr.broadcast("any/topic", b"\"v\"").await; + let record_perms = simple_record_perms(); + let mut streams: Vec<_> = (0..5) + .map(|_| mgr.subscribe("#", record_perms.clone()).1) + .collect(); + mgr.broadcast("any/topic", 0, b"\"v\"").await; for s in &mut streams { let update = s.next().await.unwrap(); assert_eq!(update.topic.as_deref(), Some("any/topic")); @@ -289,6 +324,7 @@ mod tests { #[tokio::test] async fn broadcast_prunes_closed_channel_matched_pattern() { let mgr = ClientManager::new(256); + let record_perms = simple_record_perms(); let (tx, rx) = mpsc::channel::(256); mgr.subs.insert( 1, @@ -296,17 +332,18 @@ mod tests { pattern: "dropped_channel".to_string(), tx, dropped: AtomicU64::new(0), + record_perms, }, ); drop(rx); assert_eq!(mgr.subscription_count(), 1); // Non-matching patterns survive the prune - mgr.broadcast("false_pattern", b"v").await; + mgr.broadcast("false_pattern", 0, b"v").await; assert_eq!(mgr.subscription_count(), 1); // Matched pattern get pruned - mgr.broadcast("dropped_channel", b"v").await; + mgr.broadcast("dropped_channel", 0, b"v").await; assert_eq!(mgr.subscription_count(), 0); } @@ -315,8 +352,11 @@ mod tests { #[tokio::test] async fn broadcast_shares_one_payload_to_all() { let mgr = ClientManager::new(256); - let mut streams: Vec<_> = (0..8).map(|_| mgr.subscribe("#").1).collect(); - mgr.broadcast("t", b"123").await; + let record_perms = simple_record_perms(); + let mut streams: Vec<_> = (0..8) + .map(|_| mgr.subscribe("#", record_perms.clone()).1) + .collect(); + mgr.broadcast("t", 0, b"123").await; let mut updates = Vec::new(); for s in &mut streams { updates.push(s.next().await.unwrap()); @@ -333,7 +373,8 @@ mod tests { #[tokio::test] async fn subscription_dropped_when_stream_dropped() { let mgr = ClientManager::new(256); - let (_id, stream) = mgr.subscribe("quiet.topic"); + let record_perms = simple_record_perms(); + let (_id, stream) = mgr.subscribe("quiet.topic", record_perms); // Count before stream dropping assert_eq!(mgr.subscription_count(), 1); diff --git a/aimdb-websocket-connector/src/server/connector.rs b/aimdb-websocket-connector/src/server/connector.rs index 7121a8ff..2747b51b 100644 --- a/aimdb-websocket-connector/src/server/connector.rs +++ b/aimdb-websocket-connector/src/server/connector.rs @@ -19,8 +19,8 @@ use aimdb_core::transport::{Connector, ConnectorConfig, PublishError}; use super::client_manager::ClientManager; -/// Shared late-join cache: topic → last serialized bytes. -pub(crate) type SnapshotCache = Arc>>>; +/// Shared late-join cache: (RecordId, topic) → last serialized bytes. +pub(crate) type SnapshotCache = Arc>>>; /// Outbound sink: feeds each serialized record value into the broadcast bus. pub(crate) struct WsBusSink { @@ -34,21 +34,33 @@ impl Connector for WsBusSink { fn publish( &self, destination: &str, - _config: &ConnectorConfig, + config: &ConnectorConfig, payload: &[u8], ) -> Pin> + Send + '_>> { // Own the args so the returned future borrows only `&self` (the trait // binds the future's lifetime to the receiver, not the arguments). let dest = destination.to_string(); let bytes = payload.to_vec(); + let record_index = config.record_index; Box::pin(async move { - if let Some(map) = &self.snapshot { - map.lock().unwrap().insert(dest.clone(), bytes.clone()); - } // The bus carries raw record-value bytes tagged with the topic; the // per-connection AimX codec applies the `event` envelope downstream. - self.client_mgr.broadcast(&dest, &bytes).await; - Ok(()) + if let Some(index) = record_index { + if let Some(map) = &self.snapshot { + map.lock() + .unwrap() + .insert((index, dest.clone()), bytes.clone()); + } + self.client_mgr.broadcast(&dest, index, &bytes).await; + Ok(()) + } else { + #[cfg(feature = "tracing")] + tracing::warn!( + "WsBusSink.publish for dest {} not fired due to null record_index", + &dest + ); + Err(PublishError::InvalidDestination) + } }) } } diff --git a/aimdb-websocket-connector/src/server/dispatch.rs b/aimdb-websocket-connector/src/server/dispatch.rs index ee849be9..ab5746d3 100644 --- a/aimdb-websocket-connector/src/server/dispatch.rs +++ b/aimdb-websocket-connector/src/server/dispatch.rs @@ -16,7 +16,7 @@ use std::collections::HashMap; use std::sync::Arc; -use aimdb_core::remote::{QueryHandlerFn, QueryHandlerParams, QUERY_ALL_PATTERN}; +use aimdb_core::remote::{QueryHandlerFn, QueryHandlerParams, QueryRecord, QUERY_ALL_PATTERN}; use aimdb_core::session::Session; use aimdb_core::{ AuthError, BoxFut, BoxStream, Dispatch, Payload, PeerInfo, RpcError, SessionCtx, SubUpdate, @@ -24,7 +24,7 @@ use aimdb_core::{ use serde_json::Value; use super::{ - auth::{AuthHandler, ClientId, ClientInfo, Permissions}, + auth::{AuthHandler, ClientId, ClientInfo, Permissions, RecordsBits}, client_manager::ClientManager, session::{QueryHandler, Router, SnapshotProvider}, }; @@ -69,6 +69,7 @@ impl Dispatch for WsDispatch { id: ClientId(0), remote_addr: ([0, 0, 0, 0], 0).into(), permissions: Permissions::default(), + record_perms: Arc::new(RecordsBits::new(0)), }) }); Box::new(WsSession { @@ -115,11 +116,7 @@ impl Session for WsSession { // name core can't resolve. let mut records = Vec::new(); for mut record in self.db.list_records() { - if !self - .auth - .authorize_list(&self.info, &record.record_key) - .await - { + if !self.info.record_perms.is_allowed(record.record_id as usize) { continue; } if record.schema_type.is_none() { @@ -142,17 +139,27 @@ impl Session for WsSession { }) } + /// Called when a client subscribes to a topic + /// As `Auth` no longer gate topic (permissions lie in record keys), + /// a client could subscribe to a topic which has no associated records. + /// As a result, no message could reach that client. fn subscribe<'a>( &'a mut self, topic: &'a str, ) -> BoxFut<'a, Result, RpcError>> { + // Deny subscription for no-grant clients, + // as clients having grants could have empty permission bitmap. + if self.info.permissions.read_patterns.is_empty() { + return Box::pin(async move { Err(RpcError::Denied) }); + } + Box::pin(async move { - // Per-operation authorization via the async `AuthHandler` hook. - if !self.auth.authorize_subscribe(&self.info, topic).await { - return Err(RpcError::Denied); - } // Register on the shared bus; the engine owns and drops the stream. - let (_sub_id, stream) = self.client_mgr.subscribe(topic); + // A topic always come with a record, so the client always receive a stream, + // message broadcasting will check for granted permissions (allowed records) later + let (_sub_id, stream) = self + .client_mgr + .subscribe(topic, self.info.record_perms.clone()); Ok(stream) }) } @@ -161,10 +168,13 @@ impl Session for WsSession { if !self.late_join { return Vec::new(); } + + // Filtered by client's record permissions self.snapshot_provider .snapshots(topic) .into_iter() - .map(|(topic, bytes)| (topic, Payload::from(bytes.as_slice()))) + .filter(|(record_id, _, _)| self.info.record_perms.is_allowed(*record_id)) + .map(|(_, topic, bytes)| (topic, Payload::from(bytes.as_slice()))) .collect() } @@ -188,8 +198,9 @@ impl WsSession { /// `record.query` with the shared `{name, limit, start, end}` params and /// `{records, total}` result: a plugged-in /// [`QueryHandler`] wins; otherwise delegate to the Extensions-registered - /// `QueryHandlerFn` (`with_persistence`); neither → `NotFound`. The pattern - /// passes [`AuthHandler::authorize_query`] before either is consulted. + /// `QueryHandlerFn` (`with_persistence`); neither → `NotFound`. + /// A client having no grant will have its query denied. + /// Grants allow clients to query both live and persist records. async fn record_query(&self, params: Value) -> Result { let name = params .get("name") @@ -203,12 +214,16 @@ impl WsSession { let start = params.get("start").and_then(|v| v.as_u64()); let end = params.get("end").and_then(|v| v.as_u64()); - if !self.auth.authorize_query(&self.info, &name).await { + // No-grant clients have their requests denied + // Grant clients whose perms matching persisted records, + // but matching no live records, are not denied + if self.info.permissions.read_patterns.is_empty() { return Err(RpcError::Denied); - } + }; - if let Some(handler) = &self.query_handler { - let (records, total) = handler + // The query is handled here + let records: Vec = if let Some(handler) = &self.query_handler { + let (records, _total) = handler .handle_query(&name, start, end, limit) .await .map_err(|_e| { @@ -216,26 +231,44 @@ impl WsSession { tracing::warn!("record.query handler failed: {}", _e); RpcError::Internal })?; - return Ok(serde_json::json!({ "records": records, "total": total })); - } + records + } else { + let handler_fut = { + let handler = self + .db + .extensions() + .get::() + .ok_or(RpcError::NotFound)?; + handler(QueryHandlerParams { + name, + limit, + start, + end, + }) + }; + let mut values = handler_fut.await.map_err(|_e| { + #[cfg(feature = "tracing")] + tracing::warn!("record.query persistence handler failed: {}", _e); + RpcError::Internal + })?; - let handler_fut = { - let handler = self - .db - .extensions() - .get::() - .ok_or(RpcError::NotFound)?; - handler(QueryHandlerParams { - name, - limit, - start, - end, - }) + // The QueryHandlerFn returns json, so that needs to be parsed back to Vec + // so the output could be filtered by per-client grant. + // A malformed json is QueryHandlerFn bug and surfaces as `Internal` + match values.get_mut("records").map(Value::take) { + None | Some(Value::Null) => Vec::new(), + Some(r) => serde_json::from_value::>(r).map_err(|_e| { + #[cfg(feature = "tracing")] + tracing::warn!("record.query handler returned malformed records: {}", _e); + RpcError::Internal + })?, + } }; - handler_fut.await.map_err(|_e| { - #[cfg(feature = "tracing")] - tracing::warn!("record.query persistence handler failed: {}", _e); - RpcError::Internal - }) + + // Query results need to be filtered by grant patterns, + // This ensures that clients could query reach persist records in store, + let records = self.auth.authorize_query_record(&self.info, records); + + Ok(serde_json::json!({ "records": records, "total": records.len() })) } } diff --git a/aimdb-websocket-connector/src/server/http.rs b/aimdb-websocket-connector/src/server/http.rs index b39c9ab1..a21d7660 100644 --- a/aimdb-websocket-connector/src/server/http.rs +++ b/aimdb-websocket-connector/src/server/http.rs @@ -56,7 +56,7 @@ use tower_http::cors::CorsLayer; use crate::transport::WsServerConnection; use super::{ - auth::{AuthError, AuthRequest, ClientInfo, DynAuthHandler}, + auth::{AuthError, AuthRequest, ClientInfo, DynAuthHandler, RecordsBits}, client_manager::ClientManager, }; @@ -82,6 +82,8 @@ pub(crate) struct ServerState { /// Per-connection subscription cap. pub max_subs_per_connection: usize, pub started_at: Instant, + /// List records' unique key, ordered by registration order + pub records: Arc>, } // ════════════════════════════════════════════════════════════════════ @@ -225,6 +227,12 @@ async fn ws_upgrade_handler( } }; + // Build record bits based on permissions.read_patterns + let record_perms = Arc::new(RecordsBits::resolve_permissions( + &state.records, + &permissions, + )); + // Resolve identity synchronously, before the upgrade, and carry it into the // engine via `PeerInfo::ext` (WS-style `reads_hello:false`). let id = state.client_mgr.next_client_id(); @@ -232,6 +240,7 @@ async fn ws_upgrade_handler( id, remote_addr, permissions, + record_perms, }; #[cfg(feature = "tracing")] diff --git a/aimdb-websocket-connector/src/server/session.rs b/aimdb-websocket-connector/src/server/session.rs index 1f23a0b5..d38f9e30 100644 --- a/aimdb-websocket-connector/src/server/session.rs +++ b/aimdb-websocket-connector/src/server/session.rs @@ -36,6 +36,9 @@ pub type QueryFuture<'a> = /// records. The trait is async to support database I/O. pub trait QueryHandler: Send + Sync + 'static { /// Execute a history query and return `(records, total_count)`. + /// `total_count` is limited to the handler's pattern match count, + /// and could be further restricted by clients' + /// [`Permissions::read_patterns`](crate::Permissions::read_patterns). /// /// - `pattern` — topic pattern (MQTT wildcards, `"*"` for all) /// - `from` / `to` — time range (inclusive; units are the handler's @@ -55,18 +58,22 @@ pub trait QueryHandler: Send + Sync + 'static { // ════════════════════════════════════════════════════════════════════ /// Provides the current serialized values covered by a subscription pattern for -/// late-join snapshots (one `(topic, value)` pair per covered record — a -/// wildcard pattern may cover several; an exact topic matches itself). +/// late-join snapshots (one `(record id, topic, value)` tuple per covered pair of record - topic +/// — a wildcard pattern may cover several; an exact topic may match several records). +/// The returned snapshots need record id, different clients may have different +/// read permissions to different records pub trait SnapshotProvider: Send + Sync + 'static { /// Return the latest serialized values for every topic matching `pattern`. - fn snapshots(&self, pattern: &str) -> Vec<(String, Vec)>; + /// As record keys now rule the broadcasting instead of topic, + /// return must be tuple of (record index, topic, payload) + fn snapshots(&self, pattern: &str) -> Vec<(usize, String, Vec)>; } /// A snapshot provider that always returns nothing (late-join disabled or no data). pub struct NoSnapshot; impl SnapshotProvider for NoSnapshot { - fn snapshots(&self, _pattern: &str) -> Vec<(String, Vec)> { + fn snapshots(&self, _pattern: &str) -> Vec<(usize, String, Vec)> { Vec::new() } } diff --git a/aimdb-websocket-connector/tests/decouple_record_keys_topics.rs b/aimdb-websocket-connector/tests/decouple_record_keys_topics.rs new file mode 100644 index 00000000..76edd632 --- /dev/null +++ b/aimdb-websocket-connector/tests/decouple_record_keys_topics.rs @@ -0,0 +1,1129 @@ +//! `AimDb::collect_outbound_routes` over the `ws` scheme. +//! Record keys and ws topics are decoupled (#215). +//! Grants described record keys. +//! +//! Connectors call this during `build()` to spawn one publisher task per +//! configured `link_to("ws://…")`. The returned order must track record +//! registration order, since record ids index into it. +//! +//! Several behaviors tested: per-record gating, late-join snapshots, +//! `record.list` and `record.query`, zero-grant denial, and topic-based +//! write grants; + +#![cfg(feature = "server")] + +use std::collections::HashMap; +use std::net::SocketAddr; +use std::sync::Arc; +use std::time::Duration; + +use aimdb_core::connector::TopicProvider; +use aimdb_core::remote::QueryHandlerFn; +use aimdb_core::{builder::AimDb, remote::QueryHandlerParams}; +use tokio::time::timeout; +use tokio_tungstenite::tungstenite::Message; + +use aimdb_core::buffer::BufferCfg; +use aimdb_core::AimDbBuilder; +use aimdb_tokio_adapter::{TokioAdapter, TokioRecordRegistrarExt}; +use aimdb_websocket_connector::{AuthError, AuthHandler, Permissions, WebSocketConnector}; +use serde::{Deserialize, Serialize}; +use serde_json::{json, Value}; + +use core::future::Future; +use futures_util::{SinkExt, StreamExt}; +use tokio::net::TcpStream; +use tokio_tungstenite::tungstenite::Error; + +#[derive(Debug, Clone, Serialize, Deserialize)] +struct Msg { + v: u64, +} + +type WsClient = + tokio_tungstenite::WebSocketStream>; + +async fn ws_connect(addr: SocketAddr) -> WsClient { + ws_connect_with(addr, "").await +} + +async fn ws_connect_with(addr: SocketAddr, extra: &str) -> WsClient { + // Every real client declares its AimX version at the upgrade; go through the + // shared helper so the tests exercise the exact URL the dialers produce. + let url = aimdb_core::remote::ws_url_with_version(&format!("ws://{addr}/ws")); + tokio_tungstenite::connect_async(format!("{url}{extra}")) + .await + .expect("connect") + .0 +} + +async fn try_ws_connect_with(addr: SocketAddr, extra: &str) -> Result> { + // Every real client declares its AimX version at the upgrade; go through the + // shared helper so the tests exercise the exact URL the dialers produce. + let url = aimdb_core::remote::ws_url_with_version(&format!("ws://{addr}/ws")); + let (ws, _resp) = tokio_tungstenite::connect_async(format!("{url}{extra}")).await?; + Ok(ws) +} + +/// Send one raw AimX frame (a JSON value) as a WS text message. +async fn ws_send(c: &mut WsClient, frame: Value) { + c.send(Message::Text(frame.to_string().into())) + .await + .unwrap(); +} + +/// Read the next AimX frame as JSON, with a timeout so a hang fails loudly. +async fn ws_recv(c: &mut WsClient) -> Value { + loop { + match timeout(Duration::from_secs(3), c.next()) + .await + .expect("recv timed out") + { + Some(Ok(Message::Text(t))) => return serde_json::from_str(&t).unwrap(), + Some(Ok(Message::Binary(b))) => return serde_json::from_slice(&b).unwrap(), + Some(Ok(Message::Ping(_))) | Some(Ok(Message::Pong(_))) => continue, + other => panic!("unexpected ws frame: {other:?}"), + } + } +} + +async fn try_ws_recv(c: &mut WsClient) -> Result { + loop { + match timeout(Duration::from_secs(3), c.next()) + .await + .map_err(|_| "recv timed out".to_string()) + { + Ok(Some(Ok(Message::Text(t)))) => return Ok(serde_json::from_str(&t).unwrap()), + Ok(Some(Ok(Message::Binary(b)))) => return Ok(serde_json::from_slice(&b).unwrap()), + Ok(Some(Ok(Message::Ping(_)))) | Ok(Some(Ok(Message::Pong(_)))) => continue, + Err(_) => Err("recv timed out".to_string()), + other => Err(format!("unexpected ws frame: {other:?}")), + }? + } +} + +/// Read frames until one has `"t" == tag`; panics on timeout. +async fn ws_recv_tag(c: &mut WsClient, tag: &str) -> Value { + for _ in 0..50 { + let v = ws_recv(c).await; + if v["t"] == tag { + return v; + } + } + panic!("no '{tag}' frame arrived"); +} + +/// Similar to ws_recv_tag return Err instead of panic +async fn try_ws_recv_tag(c: &mut WsClient, tag: &str) -> Result { + for _ in 0..50 { + let v = try_ws_recv(c).await?; + if v["t"] == tag { + return Ok(v); + } + } + Err("not found".to_string()) +} + +/// Wait until the server is accepting connections at `addr`. +async fn wait_for_listen(addr: SocketAddr) { + for _ in 0..200 { + if TcpStream::connect(addr).await.is_ok() { + return; + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + panic!("server never bound at {addr}"); +} + +/// Grab a probably-free ephemeral port (the WS builder binds internally and does +/// not surface `:0`'s assigned port, so we pick one up front). +fn free_addr() -> SocketAddr { + let l = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); + let a = l.local_addr().unwrap(); + drop(l); + a +} + +#[tokio::test] +async fn collect_outbound_routes_preserves_record_order() { + // A dummy address for server building only + let addr = free_addr(); + + let keys_topics = [ + ("garage", "peripheral"), + ("house", "peripheral"), + ("basement", "movement"), + ]; + + // Populate record id + let mut record_keys: HashMap<&str, usize> = HashMap::new(); + + let mut index = 0; + keys_topics.iter().for_each(|(k, _v)| { + record_keys.entry(*k).or_insert(index); + index += 1; + }); + + let mut sb = AimDbBuilder::new() + .runtime(Arc::new(TokioAdapter)) + .with_connector(WebSocketConnector::new().bind(addr).path("/ws")); + + keys_topics.iter().for_each(|(k, t)| { + sb.configure::(*k, |reg| { + reg.buffer(BufferCfg::SingleLatest) + .with_remote_access() + .link_to(format!("ws://{}", t).as_str()) + .with_serializer(|_ctx, m: &Msg| Ok(serde_json::to_vec(m).expect("serialize"))) + .finish(); + }); + }); + + let (server_db, _server_runner) = sb.build().await.expect("build server db"); + let outbound_routes = server_db.collect_outbound_routes("ws"); + assert_eq!(outbound_routes.len(), keys_topics.len()); + let outbound_iter = outbound_routes.into_iter(); + + outbound_iter + .into_iter() + .zip(keys_topics.iter()) + .for_each(|(route, (k, t))| { + assert_eq!(route.topic.as_str(), *t); + + let config_record_index: Vec = route + .config + .iter() + .filter(|(k, _v)| k.as_str() == "record_index") + .map(|(_k, v)| { + v.parse::() + .expect("failed to convert record_index to usize") + }) + .collect(); + + assert_eq!( + config_record_index.len(), + 1, + "config must have one tuple for record_index" + ); + assert_eq!( + config_record_index[0], + *record_keys.get(*k).expect("key must exist") + ); + }); +} + +//----------------- Custom AuthHandler and Permissions +struct FixedGrant(Permissions); +impl AuthHandler for FixedGrant { + fn authenticate<'a>( + &'a self, + _request: &'a aimdb_websocket_connector::AuthRequest, + ) -> std::pin::Pin> + Send + 'a>> { + let perms = self.0.clone(); + Box::pin(async move { Ok(perms) }) + } +} + +// Two or more distinct grants +struct DistinctGrant { + pub public: Permissions, + pub secret: Permissions, +} + +impl AuthHandler for DistinctGrant { + fn authenticate<'a>( + &'a self, + request: &'a aimdb_websocket_connector::AuthRequest, + ) -> std::pin::Pin> + Send + 'a>> { + // Grants are distinguished by a query param + // clients must carry this param in request + let perms_public = self.public.clone(); + let perms_secret = self.secret.clone(); + match request.query_params.get("data_type") { + Some(s) if s.as_str() == "public" => Box::pin(async move { Ok(perms_public) }), + Some(s) if s.as_str() == "secret" => Box::pin(async move { Ok(perms_secret) }), + _ => Box::pin(async move { + Err(AuthError { + message: "not authorized".to_string(), + }) + }), + } + } +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +struct Inject { + topic: String, + payload: Value, +} + +struct InjectTopic; +impl TopicProvider for InjectTopic { + fn topic(&self, value: &Inject) -> Option { + Some(value.topic.clone()) + } +} + +enum GrantType { + Fixed, + Distinct, + Injected, + WritePublic, + OldRecord, + PublicAndOld, + None, +} + +/// Different option for handle to test record.query +enum CustomQueryHandler { + All, + NoRecord, + Malformed, +} + +impl CustomQueryHandler { + pub fn get_handler_fn(self) -> QueryHandlerFn { + match self { + Self::All => { + let handler: QueryHandlerFn = Box::new(|p: QueryHandlerParams| { + Box::pin(async move { + // Tuple of key, value, timestamp + let all = [ + ("public.ledger", 1, 1), + ("secret.sensor", 2, 2), + ("old.sensor", 3, 3), + ]; + let records: Vec<_> = all + .iter() + .filter(|(key, _, _)| aimdb_core::topic_matches(&p.name, key)) + .map(|(key, payload, ts)| json!({"topic": key, "payload": payload, "ts": ts})) + .collect(); + let total = records.len(); + Ok(json!({"records": records, "total": total})) + }) + }); + handler + } + Self::NoRecord => Box::new(|_| Box::pin(async move { Ok(json!({"total": 0})) })), + Self::Malformed => { + Box::new(|_| Box::pin(async move { Ok(json!({"records": "malformed"})) })) + } + } + } +} + +/// Test fixture for several testing scenarios +async fn test_fixture( + addr: SocketAddr, + grant_type: GrantType, + query_handler: Option, +) -> AimDb { + // Different pair of record - topic + let keys_topics = [ + ("public.ledger", "public_info"), + ("public.data", "public_info"), + ("sensor", "public_info"), + ("secret.sensor", "secret_info"), + ]; + + // Different grant scenarios + let mut sb = match grant_type { + GrantType::Fixed => { + // One fixed grant for all clients + let perms = Permissions { + read_patterns: vec!["public.#".to_string()], + write_patterns: vec![], + }; + AimDbBuilder::new() + .runtime(Arc::new(TokioAdapter)) + .with_connector( + WebSocketConnector::new() + .with_auth(FixedGrant(perms)) + .bind(addr) + .path("/ws"), + ) + } + GrantType::Distinct => { + // Two type of grants for different type of clients + let perms_public = Permissions { + read_patterns: vec!["public.#".to_string()], + write_patterns: vec![], + }; + let perms_secret = Permissions { + read_patterns: vec!["secret.#".to_string()], + write_patterns: vec![], + }; + let perms = DistinctGrant { + public: perms_public, + secret: perms_secret, + }; + AimDbBuilder::new() + .runtime(Arc::new(TokioAdapter)) + .with_connector( + WebSocketConnector::new() + .with_auth(perms) + .bind(addr) + .path("/ws"), + ) + } + GrantType::Injected => { + // Grant for TopicProvider tests + let perms = Permissions { + read_patterns: vec!["injected.granted".to_string()], + write_patterns: vec![], + }; + AimDbBuilder::new() + .runtime(Arc::new(TokioAdapter)) + .with_connector( + WebSocketConnector::new() + .with_auth(FixedGrant(perms)) + .bind(addr) + .path("/ws"), + ) + } + GrantType::WritePublic => { + // Grants for client writing tests + let perms = Permissions { + read_patterns: vec!["public.#".to_string()], + write_patterns: vec!["writable_topic".to_string()], + }; + AimDbBuilder::new() + .runtime(Arc::new(TokioAdapter)) + .with_connector( + WebSocketConnector::new() + .with_auth(FixedGrant(perms)) + .bind(addr) + .path("/ws"), + ) + } + GrantType::None => { + // No permission granted + let perms = Permissions { + read_patterns: vec![], + write_patterns: vec![], + }; + AimDbBuilder::new() + .runtime(Arc::new(TokioAdapter)) + .with_connector( + WebSocketConnector::new() + .with_auth(FixedGrant(perms)) + .bind(addr) + .path("/ws"), + ) + } + GrantType::PublicAndOld => { + // Grants over both a live record and an unregistered, persisted one + let perms = Permissions { + read_patterns: vec!["public.#".to_string(), "old.#".to_string()], + write_patterns: vec![], + }; + AimDbBuilder::new() + .runtime(Arc::new(TokioAdapter)) + .with_connector( + WebSocketConnector::new() + .with_auth(FixedGrant(perms)) + .bind(addr) + .path("/ws"), + ) + } + GrantType::OldRecord => { + // Access to unregistered records but persist in store + let perms = Permissions { + read_patterns: vec!["old.#".to_string()], + write_patterns: vec![], + }; + AimDbBuilder::new() + .runtime(Arc::new(TokioAdapter)) + .with_connector( + WebSocketConnector::new() + .with_auth(FixedGrant(perms)) + .bind(addr) + .path("/ws"), + ) + } + }; + + keys_topics.iter().for_each(|(k, t)| { + sb.configure::(*k, |reg| { + reg.buffer(BufferCfg::SingleLatest) + .with_remote_access() + .link_to(format!("ws://{}", t).as_str()) + .with_serializer(|_ctx, m: &Msg| Ok(serde_json::to_vec(m).expect("serialize"))) + .finish(); + }); + }); + + // Extra keys for TopicProvider test + for key in ["injected.granted", "injected.denied"] { + sb.configure::(key, |reg| { + reg.buffer(BufferCfg::SpmcRing { capacity: 64 }) // don't coalesce successive values + .with_remote_access() + .link_to("ws://_") // placeholder; provider overrides per value + .with_topic_provider(InjectTopic) + .with_serializer(|_ctx, m: &Inject| { + Ok(serde_json::to_vec(&m.payload).expect("serialize")) + }) + .finish(); + }); + } + + // Write need different config + sb.configure::("cfg", |reg| { + reg.buffer(BufferCfg::SingleLatest) + .with_remote_access() + .link_from("ws://writable_topic") + .with_deserializer(|_ctx, d: &[u8]| { + serde_json::from_slice::(d).map_err(|e| e.to_string()) + }) + .finish(); + }); + + // For query handler test + if let Some(custom_handler) = query_handler { + let handler_fn = custom_handler.get_handler_fn(); + sb.extensions_mut().insert(handler_fn); + }; + + let (server_db, server_runner) = sb.build().await.expect("build server db"); + tokio::spawn(server_runner.run()); + + // Give the server a moment to bind before the client dials. + wait_for_listen(addr).await; + server_db +} + +#[tokio::test] +async fn authentication_by_record_keys() { + let addr = free_addr(); + let _server_db = test_fixture(addr, GrantType::Fixed, None).await; + + // Connect client + let mut client = ws_connect(addr).await; + + //------------- Testing for correct record gating + // Assert that record.list contains only "public.ledger" + ws_send( + &mut client, + json!({ + "t": "req", + "id": 1, + "method": "record.list", + "params": null, + }), + ) + .await; + let reply = ws_recv_tag(&mut client, "reply").await; + assert_eq!(reply["id"], 1); + let keys: Vec<&str> = reply["ok"] + .as_array() + .expect("record.list array") + .iter() + .map(|row| row["record_key"].as_str().unwrap()) + .collect(); + assert_eq!( + keys, + ["public.ledger", "public.data"], + "grant `public.#` must hide `sensor`" + ); +} + +#[tokio::test] +async fn client_receives_only_records_in_perms() { + let addr = free_addr(); + let server_db = test_fixture(addr, GrantType::Fixed, None).await; + + // Connect client + let mut client = ws_connect(addr).await; + // Client subscribes for topic "public_info" + ws_send( + &mut client, + json!({ + "t": "sub", + "id": 2, + "topic": "public_info", + }), + ) + .await; + + // Must wait till server registers client's subscription + let ack = ws_recv_tag(&mut client, "subscribed").await; + assert_eq!(ack["sub"], "2"); + + // Server broadcast + let _ = server_db.set_record_from_json("public.ledger", json!({"v": 1})); + + // Client receives + let ev = ws_recv_tag(&mut client, "event").await; + assert_eq!(ev["data"], json!({"v": 1})); + + // Server broadcast + let _ = server_db.set_record_from_json("sensor", json!({"v": 99})); + let _ = server_db.set_record_from_json("public.ledger", json!({"v": 2})); + + let ev = ws_recv_tag(&mut client, "event").await; + assert_eq!(ev["data"], json!({"v": 2})); +} + +#[tokio::test] +async fn late_join_client_receives_only_records_in_perms() { + let addr = free_addr(); + let server_db = test_fixture(addr, GrantType::Fixed, None).await; + + // A client as watcher, just for testing, no need in real use + let mut watcher = ws_connect(addr).await; + + // The watcher subscribes to topics, just to make sure that server properly broadcast + ws_send( + &mut watcher, + json!({ + "t": "sub", + "id": 1, + "topic": "public_info", + }), + ) + .await; + let ack = ws_recv_tag(&mut watcher, "subscribed").await; + assert_eq!(ack["sub"], "1"); + + // Server broadcast + let _ = server_db.set_record_from_json("public.ledger", json!({"v": 1})); + + // Make sure that message already broadcasted + let ev = ws_recv_tag(&mut watcher, "event").await; + assert_eq!(ev["data"], json!({"v": 1})); + + // Server broadcast + let _ = server_db.set_record_from_json("public.ledger", json!({"v": 2})); + let ev = ws_recv_tag(&mut watcher, "event").await; + assert_eq!(ev["data"], json!({"v": 2})); + + let _ = server_db.set_record_from_json("public.data", json!({"v": 3})); + let ev = ws_recv_tag(&mut watcher, "event").await; + assert_eq!(ev["data"], json!({"v": 3})); + + // Late joining client + let mut client = ws_connect(addr).await; + ws_send( + &mut client, + json!({ + "t": "sub", + "id": 1, + "topic": "public_info", + }), + ) + .await; + + // Must wait till server registers client's subscription + let ack = ws_recv_tag(&mut client, "subscribed").await; + assert_eq!(ack["sub"], "1"); + + // Getting snapshot, 2 keys so we must loop till we get last snapshot + let mut snapshot_data = Vec::new(); + loop { + let snap = ws_recv_tag(&mut client, "snap").await; + let last = snap + .get("last") + .unwrap_or(&Value::Bool(false)) + .as_bool() + .unwrap(); + assert_eq!(snap["sub"], "1"); + assert_eq!(snap["topic"], "public_info"); + snapshot_data.push(snap["data"].clone()); + if last { + break; + } + } + + // Assert that client get 2 snapshot for 2 record keys + snapshot_data.sort_by_key(|d| d["v"].as_u64()); + assert_eq!(snapshot_data, [json!({"v": 2}), json!({"v": 3})]); +} + +#[tokio::test] +async fn clients_disjoint_grants() { + let addr = free_addr(); + let server_db = test_fixture(addr, GrantType::Distinct, None).await; + + // A client with no authorization will be denied + assert!(try_ws_connect_with(addr, "&data_type=not_exist") + .await + .is_err()); + + // Different clients could have different grants + let mut client_public = ws_connect_with(addr, "&data_type=public").await; + ws_send( + &mut client_public, + json!({ + "t": "sub", + "id": 1, + "topic": "public_info", + }), + ) + .await; + let ack = ws_recv_tag(&mut client_public, "subscribed").await; + assert_eq!(ack["sub"], "1"); + + // Client with grant to secret.#, even subscribing to public_info, + // will receive nothing when server broadcasts to public.ledger + let mut client_secret = ws_connect_with(addr, "&data_type=secret").await; + ws_send( + &mut client_secret, + json!({ + "t": "sub", + "id": 2, + "topic": "public_info", + }), + ) + .await; + let ack = ws_recv_tag(&mut client_secret, "subscribed").await; + assert_eq!(ack["sub"], "2"); + + // Server broadcast to public.ledger + let _ = server_db.set_record_from_json("public.ledger", json!({"v": 1})); + + // Public client receive public.ledger + let ev = ws_recv_tag(&mut client_public, "event").await; + assert_eq!(ev["data"], json!({"v": 1})); + + // Secret client receives nothing + let ev = try_ws_recv_tag(&mut client_secret, "event").await; + assert!(ev.is_err()); + if let Err(s) = ev { + assert_eq!(s.as_str(), "recv timed out") + }; +} + +#[tokio::test] +async fn client_wildcard_subscription_receives_public_only() { + let addr = free_addr(); + let server_db = test_fixture(addr, GrantType::Fixed, None).await; + + let mut client = ws_connect(addr).await; + ws_send( + &mut client, + json!({ + "t": "sub", + "id": 1, + "topic": "#", + }), + ) + .await; + let ack = ws_recv_tag(&mut client, "subscribed").await; + assert_eq!(ack["sub"], "1"); + + let _ = server_db.set_record_from_json("public.ledger", json!({"v": 1})); + let ev = ws_recv_tag(&mut client, "event").await; + assert_eq!(ev["data"], json!({"v": 1})); + + let _ = server_db.set_record_from_json("secret.sensor", json!({"v": 2})); + let _ = server_db.set_record_from_json("public.ledger", json!({"v": 3})); + let ev = ws_recv_tag(&mut client, "event").await; + assert_eq!(ev["data"], json!({"v": 3})); + + // Late joining client also receive from public.ledger + let mut client_late = ws_connect(addr).await; + ws_send( + &mut client_late, + json!({ + "t": "sub", + "id": 2, + "topic": "#", + }), + ) + .await; + let ack = ws_recv_tag(&mut client_late, "subscribed").await; + assert_eq!(ack["sub"], "2"); + + let snap = ws_recv_tag(&mut client_late, "snap").await; + assert_eq!(snap["data"]["v"].as_u64(), Some(3)); +} + +#[tokio::test] +async fn topic_provider_injects_for_unsubscribed_client() { + let addr = free_addr(); + let server_db = test_fixture(addr, GrantType::Injected, None).await; + + let mut client = ws_connect(addr).await; + ws_send( + &mut client, + json!({ + "t": "sub", + "id": 1, + "topic": "#", + }), + ) + .await; + let ack = ws_recv_tag(&mut client, "subscribed").await; + assert_eq!(ack["sub"], "1"); + + // Client has permissions to injected.granted, topic _ is injected by "a.b" + let _ = server_db.set_record_from_json( + "injected.granted", + json!({ + "topic": "a.b", + "payload": 1 + }), + ); + let ev = ws_recv_tag(&mut client, "event").await; + assert_eq!(ev["topic"], "a.b"); + assert_eq!(ev["data"], 1); + + // Client keep receiving from record, not topic + let _ = server_db.set_record_from_json( + "injected.granted", + json!({ + "topic": "c.d", + "payload": 2 + }), + ); + let ev = ws_recv_tag(&mut client, "event").await; + assert_eq!(ev["topic"], "c.d"); + assert_eq!(ev["data"], 2); + + // Client does not received from ungranted record + let _ = server_db.set_record_from_json( + "injected.denied", + json!({ + "topic": "c.d", + "payload": 3 + }), + ); + let _ = server_db.set_record_from_json( + "injected.granted", + json!({ + "topic": "c.d", + "payload": 4 + }), + ); + let ev = ws_recv_tag(&mut client, "event").await; + assert_eq!(ev["topic"], "c.d"); + assert_eq!(ev["data"], 4); +} + +#[tokio::test] +async fn record_query_uphold_record_grant() { + let addr = free_addr(); + let _server_db = test_fixture(addr, GrantType::Fixed, Some(CustomQueryHandler::All)).await; + + let mut client = ws_connect(addr).await; + ws_send( + &mut client, + json!({ + "t": "sub", + "id": 1, + "topic": "#", + }), + ) + .await; + let ack = ws_recv_tag(&mut client, "subscribed").await; + assert_eq!(ack["sub"], "1"); + + // Query for all returned allowed records + ws_send( + &mut client, + json!({ + "t": "req", + "id": 1, + "method": "record.query", + "params": {"name": "#"}, + }), + ) + .await; + let reply = ws_recv_tag(&mut client, "reply").await; + assert_eq!(reply["id"], 1); + assert_eq!(reply["ok"]["total"], 1); + assert_eq!( + reply["ok"]["records"][0]["topic"], + "public.ledger".to_string() + ); + + // Query for not allow records returns none + ws_send( + &mut client, + json!({ + "t": "req", + "id": 2, + "method": "record.query", + "params": {"name": "secret.#"}, + }), + ) + .await; + let reply = ws_recv_tag(&mut client, "reply").await; + assert_eq!(reply["id"], 2); + assert_eq!(reply["ok"]["total"], 0); +} + +#[tokio::test] +async fn record_query_denied_for_no_grant() { + let addr = free_addr(); + let _server_db = test_fixture(addr, GrantType::None, None).await; + + let mut client = ws_connect(addr).await; + + ws_send( + &mut client, + json!({ + "t": "req", + "id": 2, + "method": "record.query", + "params": {"name": "secret.#"}, + }), + ) + .await; + let reply = ws_recv_tag(&mut client, "reply").await; + assert_eq!(reply["err"], "denied"); +} + +#[tokio::test] +async fn subscribe_denied_for_no_grant() { + let addr = free_addr(); + let _server_db = test_fixture(addr, GrantType::None, None).await; + + let mut client = ws_connect(addr).await; + ws_send( + &mut client, + json!({ + "t": "sub", + "id": 1, + "topic": "#", + }), + ) + .await; + + // The refusal is a `reply` carrying the subscribe id + let reply = ws_recv_tag(&mut client, "reply").await; + assert_eq!(reply, json!({"t": "reply", "id": 1, "err": "denied"})); +} + +#[tokio::test] +async fn subscribe_accepted_for_grant_matching_no_registered_record() { + let addr = free_addr(); + let server_db = test_fixture(addr, GrantType::OldRecord, None).await; + + // Grant `old.#` matches no registered record: the client still subscribes + let mut client = ws_connect(addr).await; + ws_send( + &mut client, + json!({ + "t": "sub", + "id": 1, + "topic": "#", + }), + ) + .await; + let ack = ws_recv_tag(&mut client, "subscribed").await; + assert_eq!(ack["sub"], "1"); + + // ...but stays silent, as no granted record ever publishes + let _ = server_db.set_record_from_json("public.ledger", json!({"v": 1})); + let ev = try_ws_recv_tag(&mut client, "event").await; + assert_eq!(ev, Err("recv timed out".to_string())); +} + +#[tokio::test] +async fn record_query_returns_unregistered_record_in_store() { + let addr = free_addr(); + let _server_db = test_fixture(addr, GrantType::OldRecord, Some(CustomQueryHandler::All)).await; + + let mut client = ws_connect(addr).await; + + ws_send( + &mut client, + json!({ + "t": "req", + "id": 1, + "method": "record.query", + "params": {"name": "old.#"}, + }), + ) + .await; + let reply = ws_recv_tag(&mut client, "reply").await; + assert_eq!(reply["id"], 1); + assert_eq!(reply["ok"]["total"], 1); + assert_eq!(reply["ok"]["records"][0]["topic"], "old.sensor"); + assert_eq!(reply["ok"]["records"][0]["payload"], 3); + + // Still cannot read + ws_send( + &mut client, + json!({ + "t": "req", + "id": 1, + "method": "record.query", + "params": {"name": "public.#"}, + }), + ) + .await; + let reply = ws_recv_tag(&mut client, "reply").await; + assert_eq!(reply["ok"]["total"], 0); +} + +#[tokio::test] +async fn record_query_returns_partial_rows_across_grants() { + let addr = free_addr(); + let _server_db = + test_fixture(addr, GrantType::PublicAndOld, Some(CustomQueryHandler::All)).await; + + let mut client = ws_connect(addr).await; + + // `#` matches all three stored rows; the grants cover two of them, + // one from a live record and one from an unregistered record + ws_send( + &mut client, + json!({ + "t": "req", + "id": 1, + "method": "record.query", + "params": {"name": "#"}, + }), + ) + .await; + let reply = ws_recv_tag(&mut client, "reply").await; + assert_eq!(reply["id"], 1); + assert_eq!(reply["ok"]["total"], 2); + let topics: Vec<&str> = reply["ok"]["records"] + .as_array() + .expect("records array") + .iter() + .map(|r| r["topic"].as_str().unwrap()) + .collect(); + assert_eq!(topics, ["public.ledger", "old.sensor"]); +} + +#[tokio::test] +async fn no_write_patterns_got_denied() { + let addr = free_addr(); + let server_db = test_fixture(addr, GrantType::Fixed, Some(CustomQueryHandler::All)).await; + + let mut client = ws_connect(addr).await; + ws_send( + &mut client, + json!({ + "t": "sub", + "id": 1, + "topic": "#", + }), + ) + .await; + let ack = ws_recv_tag(&mut client, "subscribed").await; + assert_eq!(ack["sub"], "1"); + + ws_send( + &mut client, + json!({ + "t": "write", + "topic": "cfg", + "payload": {"v": 1}, + }), + ) + .await; + + // FIFO on the one connection: the pong proves the write frame was processed. + ws_send(&mut client, json!({"t":"ping"})).await; + assert_eq!(ws_recv(&mut client).await, json!({"t":"pong"})); + tokio::time::sleep(Duration::from_millis(100)).await; + assert_eq!(server_db.try_latest_as_json("cfg"), None); +} + +#[tokio::test] +async fn write_patterns_uphold_topic() { + let addr = free_addr(); + let server_db = test_fixture(addr, GrantType::WritePublic, Some(CustomQueryHandler::All)).await; + + let mut client = ws_connect(addr).await; + ws_send( + &mut client, + json!({ + "t": "sub", + "id": 1, + "topic": "writable_topic", + }), + ) + .await; + let ack = ws_recv_tag(&mut client, "subscribed").await; + assert_eq!(ack["sub"], "1"); + + ws_send( + &mut client, + json!({ + "t": "write", + "topic": "writable_topic", + "payload": {"v": 1}, + }), + ) + .await; + + // FIFO on the one connection: the pong proves the write frame was processed. + ws_send(&mut client, json!({"t":"ping"})).await; + assert_eq!(ws_recv(&mut client).await, json!({"t":"pong"})); + tokio::time::sleep(Duration::from_millis(100)).await; + assert_eq!(server_db.try_latest_as_json("cfg"), Some(json!({"v": 1}))); +} + +#[tokio::test] +async fn query_handler_returns_empty_handled_as_empty() { + let addr = free_addr(); + let _server_db = test_fixture(addr, GrantType::Fixed, Some(CustomQueryHandler::NoRecord)).await; + + let mut client = ws_connect(addr).await; + ws_send( + &mut client, + json!({ + "t": "sub", + "id": 1, + "topic": "#", + }), + ) + .await; + let ack = ws_recv_tag(&mut client, "subscribed").await; + assert_eq!(ack["sub"], "1"); + + // Query for all returned allowed records + ws_send( + &mut client, + json!({ + "t": "req", + "id": 1, + "method": "record.query", + "params": {"name": "#"}, + }), + ) + .await; + let reply = ws_recv_tag(&mut client, "reply").await; + assert_eq!(reply["id"], 1); + assert_eq!(reply["ok"]["total"], 0); +} + +#[tokio::test] +async fn query_handler_malformed_handled_as_err() { + let addr = free_addr(); + let _server_db = + test_fixture(addr, GrantType::Fixed, Some(CustomQueryHandler::Malformed)).await; + + let mut client = ws_connect(addr).await; + ws_send( + &mut client, + json!({ + "t": "sub", + "id": 1, + "topic": "#", + }), + ) + .await; + let ack = ws_recv_tag(&mut client, "subscribed").await; + assert_eq!(ack["sub"], "1"); + + // Query for all returned allowed records + ws_send( + &mut client, + json!({ + "t": "req", + "id": 1, + "method": "record.query", + "params": {"name": "#"}, + }), + ) + .await; + let reply = ws_recv_tag(&mut client, "reply").await; + assert_eq!(reply["id"], 1); + assert_eq!(reply["err"], "internal".to_string()); +} diff --git a/aimdb-websocket-connector/tests/e2e.rs b/aimdb-websocket-connector/tests/e2e.rs index 3864e9cd..dd49ba68 100644 --- a/aimdb-websocket-connector/tests/e2e.rs +++ b/aimdb-websocket-connector/tests/e2e.rs @@ -31,8 +31,8 @@ use aimdb_data_contracts::{SchemaType, Streamable}; use aimdb_tokio_adapter::{TokioAdapter, TokioRecordRegistrarExt}; use aimdb_websocket_connector::transport::WsDialer; use aimdb_websocket_connector::{ - AuthError, AuthHandler, AuthRequest, ClientInfo, Permissions, QueryFuture, QueryHandler, - QueryRecord, WebSocketConnector, + AuthError, AuthHandler, AuthRequest, Permissions, QueryFuture, QueryHandler, QueryRecord, + WebSocketConnector, }; use futures_util::{SinkExt, StreamExt}; use serde::{Deserialize, Serialize}; @@ -89,30 +89,6 @@ impl AuthHandler for DenyAuth { } } -/// Allows the connection (allow-all permissions) but asynchronously *denies* -/// `secret/*` via `authorize_subscribe`. If the engine only consulted the static -/// permission set, `secret` would be allowed — so this proves the async hook gates. -struct AsyncTopicAuth; -impl AuthHandler for AsyncTopicAuth { - fn authenticate<'a>( - &'a self, - _request: &'a AuthRequest, - ) -> Pin> + Send + 'a>> { - Box::pin(async { Ok(Permissions::allow_all()) }) - } - fn authorize_subscribe<'a>( - &'a self, - _client: &'a ClientInfo, - topic: &'a str, - ) -> Pin + Send + 'a>> { - let denied = topic.starts_with("secret"); - Box::pin(async move { - tokio::task::yield_now().await; // simulate an async ACL lookup - !denied - }) - } -} - /// Grants everything to a client that asks via `?grant=all`, nothing to anyone /// else — one server, two very differently privileged clients. Only /// `authenticate` is overridden, so the read paths ride the trait defaults. @@ -945,26 +921,3 @@ async fn golden_wire_frames() { ws_send(&mut c, json!({"t":"ping"})).await; assert_eq!(ws_recv(&mut c).await, json!({"t":"pong"})); } - -// ── Async authorization over a real socket ─────────────────────────── - -#[tokio::test] -async fn async_authorize_subscribe_gates_despite_allow_all_permissions() { - let (addr, db) = spawn(WebSocketConnector::new().with_auth(AsyncTopicAuth)).await; - let mut c = ws_connect(addr).await; - - // Denied topic: permissions are allow-all, but the *async* hook says no. - // The refusal is a `reply` carrying the subscribe id + the 3-code error. - ws_send(&mut c, json!({"t":"sub","id":1,"topic":"secret.x"})).await; - assert_eq!( - ws_recv(&mut c).await, - json!({"t":"reply","id":1,"err":"denied"}) - ); - - // An allowed topic still works end-to-end. - ws_send(&mut c, json!({"t":"sub","id":2,"topic":"public.x"})).await; - assert_eq!(ws_recv(&mut c).await, json!({"t":"subscribed","sub":"2"})); - inject(&db, "public.x", json!(1)); - let ev = ws_recv_tag(&mut c, "event").await; - assert_eq!(ev["topic"], "public.x"); -}