PiCloud/crates/manager-core/src/group_kv_service.rs

//! `GroupKvServiceImpl` — wires `GroupKvRepo` + the group-collection registry
//! underneath the `picloud_shared::GroupKvService` trait that scripts reach via
//! the `kv::shared_collection("name")` Rhai handle (§11.6).
//!
//! Layers added over the raw repo:
//!
//! 1. Empty-collection rejection.
//! 2. **Owner resolution** (the structural isolation boundary): the collection
//!    name is resolved to the nearest ancestor group that declares it shared,
//!    walking the calling app's chain. No declaration on the chain → a clean
//!    `CollectionNotShared`, BEFORE any authz or storage access.
//! 3. **Reads-open / writes-authed authz**: reads use `script_gate`
//!    (anonymous public scripts skip); writes use `script_gate_require_principal`
//!    (anonymous fails closed — shared mutation always needs an authenticated
//!    editor+ on the owning group).
//! 4. Per-value size cap (reuses `PICLOUD_KV_MAX_VALUE_BYTES`) + a per-group
//!    row quota (§11.6 M3, `PICLOUD_GROUP_KV_MAX_ROWS`).
//! 5. §11.6 M2 shared-collection triggers: a write fires `shared = true`
//!    triggers on the owning group, under the writer app (`emit_shared`).

use std::sync::Arc;

use async_trait::async_trait;
use picloud_shared::{
    GroupId, GroupKvError, GroupKvService, KvListPage, NoopEventEmitter, SdkCallCx, ServiceEvent,
    ServiceEventEmitter,
};

use crate::authz::{self, AuthzRepo, Capability};
use crate::group_collection_repo::GroupCollectionResolver;
use crate::group_kv_repo::{GroupKvRepo, GroupKvRepoError};
use crate::kv_service::kv_max_value_bytes_from_env;

/// The registry `kind` this service resolves.
const KIND_KV: &str = "kv";

pub struct GroupKvServiceImpl {
    repo: Arc<dyn GroupKvRepo>,
    resolver: Arc<dyn GroupCollectionResolver>,
    authz: Arc<dyn AuthzRepo>,
    max_value_bytes: usize,
    /// §11.6 per-group quota: max total keys across the group's shared-KV
    /// collections (`PICLOUD_GROUP_KV_MAX_ROWS`). Checked only when inserting a
    /// NEW key.
    max_rows: u64,
    /// §11.6 M4 per-group quota: max total stored bytes across the group's
    /// shared-KV collections (`PICLOUD_GROUP_KV_MAX_TOTAL_BYTES`). Checked on the
    /// projected total (old value subtracted, new added) so a same/smaller update
    /// near the cap is still allowed.
    max_total_bytes: u64,
    /// §11.6: fires `shared = true` triggers on a shared-collection write.
    /// Defaults to the noop emitter; the host wires the outbox emitter via
    /// [`Self::with_events`].
    events: Arc<dyn ServiceEventEmitter>,
}

impl GroupKvServiceImpl {
    #[must_use]
    pub fn new(
        repo: Arc<dyn GroupKvRepo>,
        resolver: Arc<dyn GroupCollectionResolver>,
        authz: Arc<dyn AuthzRepo>,
    ) -> Self {
        Self::with_max_value_bytes(repo, resolver, authz, kv_max_value_bytes_from_env())
    }

    /// Wire the event emitter (§11.6 shared-collection triggers). Without this
    /// the service uses the noop emitter and shared writes fire nothing.
    #[must_use]
    pub fn with_events(mut self, events: Arc<dyn ServiceEventEmitter>) -> Self {
        self.events = events;
        self
    }

    /// Override the per-group row quota (§11.6). Mainly for tests; the host uses
    /// the env-var default.
    #[must_use]
    pub fn with_max_rows(mut self, max_rows: u64) -> Self {
        self.max_rows = max_rows;
        self
    }

    #[must_use]
    pub fn with_max_value_bytes(
        repo: Arc<dyn GroupKvRepo>,
        resolver: Arc<dyn GroupCollectionResolver>,
        authz: Arc<dyn AuthzRepo>,
        max_value_bytes: usize,
    ) -> Self {
        Self {
            events: Arc::new(NoopEventEmitter),
            max_rows: crate::group_quota::group_kv_max_rows_from_env(),
            max_total_bytes: crate::group_quota::group_kv_max_total_bytes_from_env(),
            repo,
            resolver,
            authz,
            max_value_bytes,
        }
    }

    /// Override the per-group total-bytes quota (§11.6 M4). Mainly for tests.
    #[must_use]
    pub fn with_max_total_bytes(mut self, max_total_bytes: u64) -> Self {
        self.max_total_bytes = max_total_bytes;
        self
    }

    /// The structural boundary: resolve the collection to its owning group on
    /// the calling app's chain. `CollectionNotShared` when no ancestor group
    /// declares it — the same outcome a foreign app gets, by construction.
    async fn owning_group(
        &self,
        cx: &SdkCallCx,
        collection: &str,
    ) -> Result<GroupId, GroupKvError> {
        if collection.is_empty() {
            return Err(GroupKvError::InvalidCollection);
        }
        self.resolver
            .resolve_owning_group(cx.app_id, collection, KIND_KV)
            .await
            .map_err(|e| GroupKvError::Backend(e.to_string()))?
            .ok_or_else(|| GroupKvError::CollectionNotShared(collection.to_string()))
    }

    /// §11.6: fire `shared = true` kv triggers on the owning group. Best-effort
    /// — an emit failure is logged, never surfaced (the write already
    /// committed), matching the per-app `KvServiceImpl` behaviour.
    async fn emit_shared(
        &self,
        cx: &SdkCallCx,
        group_id: GroupId,
        op: &'static str,
        collection: &str,
        key: &str,
        payload: Option<serde_json::Value>,
    ) {
        if let Err(e) = self
            .events
            .emit_shared(
                cx,
                group_id,
                ServiceEvent {
                    source: "kv",
                    op,
                    collection: Some(collection.to_string()),
                    key: Some(key.to_string()),
                    payload,
                    old_payload: None,
                },
            )
            .await
        {
            tracing::error!(error = %e, source = "kv", op, event_emit_failure = true, "shared event emit failed");
        }
    }

    async fn check_read(&self, cx: &SdkCallCx, group_id: GroupId) -> Result<(), GroupKvError> {
        authz::script_gate(
            &*self.authz,
            cx,
            Capability::GroupKvRead(group_id),
            || GroupKvError::Forbidden,
            GroupKvError::Backend,
        )
        .await
    }

    async fn check_write(&self, cx: &SdkCallCx, group_id: GroupId) -> Result<(), GroupKvError> {
        // Fails closed for an anonymous principal: shared mutation always needs
        // an authenticated editor+ on the owning group.
        authz::script_gate_require_principal(
            &*self.authz,
            cx,
            Capability::GroupKvWrite(group_id),
            || GroupKvError::Forbidden,
            GroupKvError::Backend,
        )
        .await
    }
}

impl From<GroupKvRepoError> for GroupKvError {
    fn from(e: GroupKvRepoError) -> Self {
        Self::Backend(e.to_string())
    }
}

#[async_trait]
impl GroupKvService for GroupKvServiceImpl {
    async fn get(
        &self,
        cx: &SdkCallCx,
        collection: &str,
        key: &str,
    ) -> Result<Option<serde_json::Value>, GroupKvError> {
        let group_id = self.owning_group(cx, collection).await?;
        self.check_read(cx, group_id).await?;
        Ok(self.repo.get(group_id, collection, key).await?)
    }

    async fn set(
        &self,
        cx: &SdkCallCx,
        collection: &str,
        key: &str,
        value: serde_json::Value,
    ) -> Result<(), GroupKvError> {
        let group_id = self.owning_group(cx, collection).await?;
        let encoded_len = serde_json::to_vec(&value)
            .map(|v| v.len())
            .map_err(|e| GroupKvError::Backend(format!("encode value: {e}")))?;
        if encoded_len > self.max_value_bytes {
            return Err(GroupKvError::ValueTooLarge {
                limit: self.max_value_bytes,
                actual: encoded_len,
            });
        }
        self.check_write(cx, group_id).await?;
        // Determine insert vs update for the event op + quota deltas (best-effort,
        // like the per-app path — the get+set is non-transactional but triggers
        // are fire-and-forget and quotas are safety rails, not exact accounting).
        let existing = self.repo.get(group_id, collection, key).await?;
        let existed = existing.is_some();
        // §11.6 row quota: a NEW key must fit under the group's row ceiling; an
        // update of an existing key is net-zero and exempt.
        if !existed {
            let count = self.repo.count_rows(group_id).await?;
            if count >= self.max_rows {
                return Err(GroupKvError::QuotaExceeded {
                    limit: usize::try_from(self.max_rows).unwrap_or(usize::MAX),
                    actual: usize::try_from(count).unwrap_or(usize::MAX),
                });
            }
        }
        // §11.6 M4 byte quota: the PROJECTED total (old value's bytes subtracted,
        // new value's added) must fit under the group's total-bytes ceiling, so a
        // same-or-smaller update near the cap is still allowed.
        let old_len = existing
            .as_ref()
            .and_then(|v| serde_json::to_vec(v).ok())
            .map_or(0u64, |b| b.len() as u64);
        let used = self.repo.total_bytes(group_id).await?;
        let projected = used.saturating_sub(old_len) + encoded_len as u64;
        if projected > self.max_total_bytes {
            return Err(GroupKvError::TotalBytesQuotaExceeded {
                limit: self.max_total_bytes,
                actual: projected,
            });
        }
        self.repo
            .set(group_id, collection, key, value.clone())
            .await?;
        self.emit_shared(
            cx,
            group_id,
            if existed { "update" } else { "insert" },
            collection,
            key,
            Some(value),
        )
        .await;
        Ok(())
    }

    async fn delete(
        &self,
        cx: &SdkCallCx,
        collection: &str,
        key: &str,
    ) -> Result<bool, GroupKvError> {
        let group_id = self.owning_group(cx, collection).await?;
        self.check_write(cx, group_id).await?;
        let deleted = self.repo.delete(group_id, collection, key).await?.is_some();
        if deleted {
            self.emit_shared(cx, group_id, "delete", collection, key, None)
                .await;
        }
        Ok(deleted)
    }

    async fn has(&self, cx: &SdkCallCx, collection: &str, key: &str) -> Result<bool, GroupKvError> {
        let group_id = self.owning_group(cx, collection).await?;
        self.check_read(cx, group_id).await?;
        Ok(self.repo.has(group_id, collection, key).await?)
    }

    async fn list(
        &self,
        cx: &SdkCallCx,
        collection: &str,
        cursor: Option<&str>,
        limit: u32,
    ) -> Result<KvListPage, GroupKvError> {
        let group_id = self.owning_group(cx, collection).await?;
        self.check_read(cx, group_id).await?;
        Ok(self.repo.list(group_id, collection, cursor, limit).await?)
    }
}

// ----------------------------------------------------------------------------
// Tests — in-memory repo + a fake resolver so unit tests don't need Postgres.
// ----------------------------------------------------------------------------

#[cfg(test)]
mod tests {
    use super::*;
    use crate::authz::{AuthzError, AuthzRepo};
    use picloud_shared::{
        AdminUserId, AppId, AppRole, ExecutionId, InstanceRole, Principal, RequestId, ScriptId,
        UserId,
    };
    use std::collections::{BTreeMap, HashMap};
    use tokio::sync::Mutex;

    #[derive(Default)]
    struct InMemoryGroupKvRepo {
        data: Mutex<BTreeMap<(GroupId, String, String), serde_json::Value>>,
    }

    #[async_trait]
    impl GroupKvRepo for InMemoryGroupKvRepo {
        async fn get(
            &self,
            group_id: GroupId,
            collection: &str,
            key: &str,
        ) -> Result<Option<serde_json::Value>, GroupKvRepoError> {
            Ok(self
                .data
                .lock()
                .await
                .get(&(group_id, collection.to_string(), key.to_string()))
                .cloned())
        }

        async fn set(
            &self,
            group_id: GroupId,
            collection: &str,
            key: &str,
            value: serde_json::Value,
        ) -> Result<Option<serde_json::Value>, GroupKvRepoError> {
            Ok(self
                .data
                .lock()
                .await
                .insert((group_id, collection.to_string(), key.to_string()), value))
        }

        async fn delete(
            &self,
            group_id: GroupId,
            collection: &str,
            key: &str,
        ) -> Result<Option<serde_json::Value>, GroupKvRepoError> {
            Ok(self
                .data
                .lock()
                .await
                .remove(&(group_id, collection.to_string(), key.to_string())))
        }

        async fn has(
            &self,
            group_id: GroupId,
            collection: &str,
            key: &str,
        ) -> Result<bool, GroupKvRepoError> {
            Ok(self.data.lock().await.contains_key(&(
                group_id,
                collection.to_string(),
                key.to_string(),
            )))
        }

        async fn list(
            &self,
            group_id: GroupId,
            collection: &str,
            _cursor: Option<&str>,
            limit: u32,
        ) -> Result<KvListPage, GroupKvRepoError> {
            let data = self.data.lock().await;
            let mut keys: Vec<String> = data
                .iter()
                .filter(|((g, c, _), _)| *g == group_id && c == collection)
                .map(|((_, _, k), _)| k.clone())
                .collect();
            keys.sort();
            keys.truncate((limit as usize).max(1));
            Ok(KvListPage {
                keys,
                next_cursor: None,
            })
        }

        async fn count_rows(&self, group_id: GroupId) -> Result<u64, GroupKvRepoError> {
            let data = self.data.lock().await;
            Ok(data.keys().filter(|(g, _, _)| *g == group_id).count() as u64)
        }

        async fn total_bytes(&self, group_id: GroupId) -> Result<u64, GroupKvRepoError> {
            let data = self.data.lock().await;
            Ok(data
                .iter()
                .filter(|((g, _, _), _)| *g == group_id)
                .map(|(_, v)| serde_json::to_vec(v).map_or(0, |b| b.len() as u64))
                .sum())
        }
    }

    /// Maps `(app_id, lowercased name) -> owning group`. Anything absent is
    /// "not shared with you".
    #[derive(Default)]
    struct FakeResolver {
        map: HashMap<(AppId, String), GroupId>,
    }

    #[async_trait]
    impl GroupCollectionResolver for FakeResolver {
        async fn resolve_owning_group(
            &self,
            app_id: AppId,
            name: &str,
            _kind: &str,
        ) -> Result<Option<GroupId>, sqlx::Error> {
            Ok(self.map.get(&(app_id, name.to_lowercase())).copied())
        }
    }

    #[derive(Default)]
    struct DenyingAuthzRepo;

    #[async_trait]
    impl AuthzRepo for DenyingAuthzRepo {
        async fn membership(
            &self,
            _user_id: UserId,
            _app_id: AppId,
        ) -> Result<Option<AppRole>, AuthzError> {
            Ok(None)
        }
    }

    fn cx_with(app_id: AppId, principal: Option<Principal>) -> SdkCallCx {
        SdkCallCx {
            app_id,
            script_id: ScriptId::new(),
            principal,
            execution_id: ExecutionId::new(),
            request_id: RequestId::new(),
            trigger_depth: 0,
            root_execution_id: ExecutionId::new(),
            is_dead_letter_handler: false,
            event: None,
        }
    }

    fn owner() -> Principal {
        Principal {
            user_id: AdminUserId::new(),
            instance_role: InstanceRole::Owner,
            scopes: None,
            app_binding: None,
        }
    }

    fn member_no_role() -> Principal {
        Principal {
            user_id: AdminUserId::new(),
            instance_role: InstanceRole::Member,
            scopes: None,
            app_binding: None,
        }
    }

    fn svc(resolver: FakeResolver) -> GroupKvServiceImpl {
        GroupKvServiceImpl::new(
            Arc::new(InMemoryGroupKvRepo::default()),
            Arc::new(resolver),
            Arc::new(DenyingAuthzRepo),
        )
    }

    #[tokio::test]
    async fn per_group_row_quota_rejects_new_keys_but_allows_updates() {
        // §11.6: a group's shared-KV store has a row ceiling; a NEW key over the
        // cap is rejected, but updating an existing key (net-zero) still works.
        let app = AppId::new();
        let group = GroupId::new();
        let mut resolver = FakeResolver::default();
        resolver.map.insert((app, "catalog".into()), group);
        let kv = GroupKvServiceImpl::new(
            Arc::new(InMemoryGroupKvRepo::default()),
            Arc::new(resolver),
            Arc::new(DenyingAuthzRepo),
        )
        .with_max_rows(2);
        let cx = cx_with(app, Some(owner()));

        kv.set(&cx, "catalog", "a", serde_json::json!(1))
            .await
            .unwrap();
        kv.set(&cx, "catalog", "b", serde_json::json!(2))
            .await
            .unwrap();
        // A third distinct key exceeds the cap.
        let err = kv
            .set(&cx, "catalog", "c", serde_json::json!(3))
            .await
            .unwrap_err();
        assert!(
            matches!(
                err,
                GroupKvError::QuotaExceeded {
                    limit: 2,
                    actual: 2
                }
            ),
            "got {err:?}"
        );
        // Updating an existing key is exempt (net-zero).
        kv.set(&cx, "catalog", "a", serde_json::json!(11))
            .await
            .expect("update of an existing key must be allowed at the cap");
    }

    #[tokio::test]
    async fn per_group_byte_quota_uses_the_projected_total() {
        // §11.6 M4: the byte ceiling is checked on the PROJECTED total, so a
        // write that would push the group over is rejected, but a same/smaller
        // update near the cap is still allowed (unlike a naive used+new check).
        let app = AppId::new();
        let group = GroupId::new();
        let mut resolver = FakeResolver::default();
        resolver.map.insert((app, "catalog".into()), group);
        // A 20-byte ceiling; each `"xxxxx"` value encodes to 7 bytes ("\"xxxxx\"").
        let kv = GroupKvServiceImpl::new(
            Arc::new(InMemoryGroupKvRepo::default()),
            Arc::new(resolver),
            Arc::new(DenyingAuthzRepo),
        )
        .with_max_rows(1000)
        .with_max_total_bytes(20);
        let cx = cx_with(app, Some(owner()));

        kv.set(&cx, "catalog", "a", serde_json::json!("xxxxx"))
            .await
            .unwrap(); // used = 7
        kv.set(&cx, "catalog", "b", serde_json::json!("xxxxx"))
            .await
            .unwrap(); // used = 14
                       // A third 7-byte value → projected 21 > 20 → rejected.
        let err = kv
            .set(&cx, "catalog", "c", serde_json::json!("xxxxx"))
            .await
            .unwrap_err();
        assert!(
            matches!(err, GroupKvError::TotalBytesQuotaExceeded { limit: 20, .. }),
            "got {err:?}"
        );
        // A same-size update of an existing key stays within budget (projected
        // = 14 - 7 + 7 = 14 ≤ 20) — the delta rule, not used+new.
        kv.set(&cx, "catalog", "a", serde_json::json!("yyyyy"))
            .await
            .expect("a same-size update near the cap must be allowed");
        // Growing an existing value past the cap is still rejected.
        let err = kv
            .set(&cx, "catalog", "a", serde_json::json!("zzzzzzzzzzzzzzz"))
            .await
            .unwrap_err();
        assert!(
            matches!(err, GroupKvError::TotalBytesQuotaExceeded { .. }),
            "got {err:?}"
        );
    }

    #[tokio::test]
    async fn unrelated_app_gets_collection_not_shared() {
        // app_a's chain declares "catalog"; app_b's does not.
        let app_a = AppId::new();
        let app_b = AppId::new();
        let group = GroupId::new();
        let mut resolver = FakeResolver::default();
        resolver.map.insert((app_a, "catalog".into()), group);
        let kv = svc(resolver);

        // app_a (owner principal) can write+read.
        let cx_a = cx_with(app_a, Some(owner()));
        kv.set(&cx_a, "catalog", "k", serde_json::json!(1))
            .await
            .unwrap();
        assert_eq!(
            kv.get(&cx_a, "catalog", "k").await.unwrap(),
            Some(serde_json::json!(1))
        );

        // app_b is off-chain — the name does not resolve.
        let cx_b = cx_with(app_b, Some(owner()));
        let err = kv.get(&cx_b, "catalog", "k").await.unwrap_err();
        assert!(matches!(err, GroupKvError::CollectionNotShared(c) if c == "catalog"));
    }

    #[tokio::test]
    async fn reads_open_writes_require_auth() {
        let app = AppId::new();
        let group = GroupId::new();
        let mut resolver = FakeResolver::default();
        resolver.map.insert((app, "catalog".into()), group);
        let kv = svc(resolver);

        // Seed a value as owner.
        let owner_cx = cx_with(app, Some(owner()));
        kv.set(&owner_cx, "catalog", "k", serde_json::json!("v"))
            .await
            .unwrap();

        // Anonymous (principal=None) READ is allowed (reads-open).
        let anon = cx_with(app, None);
        assert_eq!(
            kv.get(&anon, "catalog", "k").await.unwrap(),
            Some(serde_json::json!("v"))
        );

        // Anonymous WRITE fails closed.
        let err = kv
            .set(&anon, "catalog", "k", serde_json::json!("x"))
            .await
            .unwrap_err();
        assert!(matches!(err, GroupKvError::Forbidden));

        // Authenticated member with no group role: write forbidden.
        let member = cx_with(app, Some(member_no_role()));
        let err = kv.delete(&member, "catalog", "k").await.unwrap_err();
        assert!(matches!(err, GroupKvError::Forbidden));
    }

    #[tokio::test]
    async fn empty_collection_rejected_before_resolve() {
        let kv = svc(FakeResolver::default());
        let cx = cx_with(AppId::new(), None);
        let err = kv.get(&cx, "", "k").await.unwrap_err();
        assert!(matches!(err, GroupKvError::InvalidCollection));
    }
}