diff --git a/src/declarations/satellite/satellite.did.d.ts b/src/declarations/satellite/satellite.did.d.ts index 369966136..907019278 100644 --- a/src/declarations/satellite/satellite.did.d.ts +++ b/src/declarations/satellite/satellite.did.d.ts @@ -191,6 +191,9 @@ export interface SetRule { rate_config: [] | [RateConfig]; write: Permission; } +export interface SetUserUsage { + items_count: number; +} export interface StorageConfig { iframe: [] | [StorageConfigIFrame]; rewrites: Array<[string, string]>; @@ -237,6 +240,12 @@ export interface UploadChunk { export interface UploadChunkResult { chunk_id: bigint; } +export interface UserUsage { + updated_at: bigint; + created_at: bigint; + version: [] | [bigint]; + items_count: number; +} export interface _SERVICE { build_version: ActorMethod<[], string>; commit_asset_upload: ActorMethod<[CommitBatch], undefined>; @@ -265,6 +274,7 @@ export interface _SERVICE { get_many_docs: ActorMethod<[Array<[string, string]>], Array<[string, [] | [Doc]]>>; get_rule: ActorMethod<[CollectionType, string], [] | [Rule]>; get_storage_config: ActorMethod<[], StorageConfig>; + get_user_usage: ActorMethod<[string, CollectionType, [] | [Principal]], [] | [UserUsage]>; http_request: ActorMethod<[HttpRequest], HttpResponse>; http_request_streaming_callback: ActorMethod< [StreamingCallbackToken], @@ -285,6 +295,7 @@ export interface _SERVICE { set_many_docs: ActorMethod<[Array<[string, string, SetDoc]>], Array<[string, Doc]>>; set_rule: ActorMethod<[CollectionType, string, SetRule], Rule>; set_storage_config: ActorMethod<[StorageConfig], undefined>; + set_user_usage: ActorMethod<[string, CollectionType, Principal, SetUserUsage], UserUsage>; upload_asset_chunk: ActorMethod<[UploadChunk], UploadChunkResult>; version: ActorMethod<[], string>; } diff --git a/src/declarations/satellite/satellite.factory.certified.did.js b/src/declarations/satellite/satellite.factory.certified.did.js index 94aaaa4fa..6233cfabf 100644 --- a/src/declarations/satellite/satellite.factory.certified.did.js +++ b/src/declarations/satellite/satellite.factory.certified.did.js @@ -145,6 +145,12 @@ export const idlFactory = ({ IDL }) => { rate_config: IDL.Opt(RateConfig), write: Permission }); + const UserUsage = IDL.Record({ + updated_at: IDL.Nat64, + created_at: IDL.Nat64, + version: IDL.Opt(IDL.Nat64), + items_count: IDL.Nat32 + }); const HttpRequest = IDL.Record({ url: IDL.Text, method: IDL.Text, @@ -231,6 +237,7 @@ export const idlFactory = ({ IDL }) => { rate_config: IDL.Opt(RateConfig), write: Permission }); + const SetUserUsage = IDL.Record({ items_count: IDL.Nat32 }); const UploadChunk = IDL.Record({ content: IDL.Vec(IDL.Nat8), batch_id: IDL.Nat, @@ -277,6 +284,11 @@ export const idlFactory = ({ IDL }) => { ), get_rule: IDL.Func([CollectionType, IDL.Text], [IDL.Opt(Rule)], []), get_storage_config: IDL.Func([], [StorageConfig], []), + get_user_usage: IDL.Func( + [IDL.Text, CollectionType, IDL.Opt(IDL.Principal)], + [IDL.Opt(UserUsage)], + [] + ), http_request: IDL.Func([HttpRequest], [HttpResponse], []), http_request_streaming_callback: IDL.Func( [StreamingCallbackToken], @@ -306,6 +318,11 @@ export const idlFactory = ({ IDL }) => { ), set_rule: IDL.Func([CollectionType, IDL.Text, SetRule], [Rule], []), set_storage_config: IDL.Func([StorageConfig], [], []), + set_user_usage: IDL.Func( + [IDL.Text, CollectionType, IDL.Principal, SetUserUsage], + [UserUsage], + [] + ), upload_asset_chunk: IDL.Func([UploadChunk], [UploadChunkResult], []), version: IDL.Func([], [IDL.Text], []) }); diff --git a/src/declarations/satellite/satellite.factory.did.js b/src/declarations/satellite/satellite.factory.did.js index 7ea9f6a8e..73de6d16f 100644 --- a/src/declarations/satellite/satellite.factory.did.js +++ b/src/declarations/satellite/satellite.factory.did.js @@ -145,6 +145,12 @@ export const idlFactory = ({ IDL }) => { rate_config: IDL.Opt(RateConfig), write: Permission }); + const UserUsage = IDL.Record({ + updated_at: IDL.Nat64, + created_at: IDL.Nat64, + version: IDL.Opt(IDL.Nat64), + items_count: IDL.Nat32 + }); const HttpRequest = IDL.Record({ url: IDL.Text, method: IDL.Text, @@ -231,6 +237,7 @@ export const idlFactory = ({ IDL }) => { rate_config: IDL.Opt(RateConfig), write: Permission }); + const SetUserUsage = IDL.Record({ items_count: IDL.Nat32 }); const UploadChunk = IDL.Record({ content: IDL.Vec(IDL.Nat8), batch_id: IDL.Nat, @@ -277,6 +284,11 @@ export const idlFactory = ({ IDL }) => { ), get_rule: IDL.Func([CollectionType, IDL.Text], [IDL.Opt(Rule)], ['query']), get_storage_config: IDL.Func([], [StorageConfig], ['query']), + get_user_usage: IDL.Func( + [IDL.Text, CollectionType, IDL.Opt(IDL.Principal)], + [IDL.Opt(UserUsage)], + ['query'] + ), http_request: IDL.Func([HttpRequest], [HttpResponse], ['query']), http_request_streaming_callback: IDL.Func( [StreamingCallbackToken], @@ -306,6 +318,11 @@ export const idlFactory = ({ IDL }) => { ), set_rule: IDL.Func([CollectionType, IDL.Text, SetRule], [Rule], []), set_storage_config: IDL.Func([StorageConfig], [], []), + set_user_usage: IDL.Func( + [IDL.Text, CollectionType, IDL.Principal, SetUserUsage], + [UserUsage], + [] + ), upload_asset_chunk: IDL.Func([UploadChunk], [UploadChunkResult], []), version: IDL.Func([], [IDL.Text], ['query']) }); diff --git a/src/declarations/satellite/satellite.factory.did.mjs b/src/declarations/satellite/satellite.factory.did.mjs index 7ea9f6a8e..73de6d16f 100644 --- a/src/declarations/satellite/satellite.factory.did.mjs +++ b/src/declarations/satellite/satellite.factory.did.mjs @@ -145,6 +145,12 @@ export const idlFactory = ({ IDL }) => { rate_config: IDL.Opt(RateConfig), write: Permission }); + const UserUsage = IDL.Record({ + updated_at: IDL.Nat64, + created_at: IDL.Nat64, + version: IDL.Opt(IDL.Nat64), + items_count: IDL.Nat32 + }); const HttpRequest = IDL.Record({ url: IDL.Text, method: IDL.Text, @@ -231,6 +237,7 @@ export const idlFactory = ({ IDL }) => { rate_config: IDL.Opt(RateConfig), write: Permission }); + const SetUserUsage = IDL.Record({ items_count: IDL.Nat32 }); const UploadChunk = IDL.Record({ content: IDL.Vec(IDL.Nat8), batch_id: IDL.Nat, @@ -277,6 +284,11 @@ export const idlFactory = ({ IDL }) => { ), get_rule: IDL.Func([CollectionType, IDL.Text], [IDL.Opt(Rule)], ['query']), get_storage_config: IDL.Func([], [StorageConfig], ['query']), + get_user_usage: IDL.Func( + [IDL.Text, CollectionType, IDL.Opt(IDL.Principal)], + [IDL.Opt(UserUsage)], + ['query'] + ), http_request: IDL.Func([HttpRequest], [HttpResponse], ['query']), http_request_streaming_callback: IDL.Func( [StreamingCallbackToken], @@ -306,6 +318,11 @@ export const idlFactory = ({ IDL }) => { ), set_rule: IDL.Func([CollectionType, IDL.Text, SetRule], [Rule], []), set_storage_config: IDL.Func([StorageConfig], [], []), + set_user_usage: IDL.Func( + [IDL.Text, CollectionType, IDL.Principal, SetUserUsage], + [UserUsage], + [] + ), upload_asset_chunk: IDL.Func([UploadChunk], [UploadChunkResult], []), version: IDL.Func([], [IDL.Text], ['query']) }); diff --git a/src/libs/collections/src/constants.rs b/src/libs/collections/src/constants.rs index 39f51f332..848936f1f 100644 --- a/src/libs/collections/src/constants.rs +++ b/src/libs/collections/src/constants.rs @@ -35,6 +35,8 @@ pub const DEFAULT_DB_COLLECTIONS: [(&str, SetRule); 2] = [ (LOG_COLLECTION_KEY, DEFAULT_DB_LOG_RULE), ]; +pub const DB_COLLECTIONS_NO_USER_USAGE: [&str; 1] = [LOG_COLLECTION_KEY]; + pub const ASSET_COLLECTION_KEY: &str = "#dapp"; pub const DEFAULT_ASSETS_COLLECTIONS: [(&str, SetRule); 1] = [( @@ -50,3 +52,5 @@ pub const DEFAULT_ASSETS_COLLECTIONS: [(&str, SetRule); 1] = [( rate_config: None, }, )]; + +pub const ASSETS_COLLECTIONS_NO_USER_USAGE: [&str; 1] = [ASSET_COLLECTION_KEY]; diff --git a/src/libs/satellite/satellite.did b/src/libs/satellite/satellite.did index a3d162e11..ebf0d5238 100644 --- a/src/libs/satellite/satellite.did +++ b/src/libs/satellite/satellite.did @@ -155,6 +155,7 @@ type SetRule = record { rate_config : opt RateConfig; write : Permission; }; +type SetUserUsage = record { items_count : nat32 }; type StorageConfig = record { iframe : opt StorageConfigIFrame; rewrites : vec record { text; text }; @@ -197,6 +198,12 @@ type UploadChunk = record { order_id : opt nat; }; type UploadChunkResult = record { chunk_id : nat }; +type UserUsage = record { + updated_at : nat64; + created_at : nat64; + version : opt nat64; + items_count : nat32; +}; service : () -> { commit_asset_upload : (CommitBatch) -> (); count_assets : (text, ListParams) -> (nat64) query; @@ -230,6 +237,9 @@ service : () -> { ) query; get_rule : (CollectionType, text) -> (opt Rule) query; get_storage_config : () -> (StorageConfig) query; + get_user_usage : (text, CollectionType, opt principal) -> ( + opt UserUsage, + ) query; http_request : (HttpRequest) -> (HttpResponse) query; http_request_streaming_callback : (StreamingCallbackToken) -> ( StreamingCallbackHttpResponse, @@ -253,6 +263,9 @@ service : () -> { ); set_rule : (CollectionType, text, SetRule) -> (Rule); set_storage_config : (StorageConfig) -> (); + set_user_usage : (text, CollectionType, principal, SetUserUsage) -> ( + UserUsage, + ); upload_asset_chunk : (UploadChunk) -> (UploadChunkResult); version : () -> (text) query; } diff --git a/src/libs/satellite/src/lib.rs b/src/libs/satellite/src/lib.rs index 4e1c97bae..3010af824 100644 --- a/src/libs/satellite/src/lib.rs +++ b/src/libs/satellite/src/lib.rs @@ -13,12 +13,16 @@ mod rules; mod satellite; mod storage; mod types; +mod usage; mod version; use crate::auth::types::config::AuthenticationConfig; use crate::db::types::config::DbConfig; use crate::guards::{caller_is_admin_controller, caller_is_controller}; -use crate::types::interface::{Config, CollectionType}; +use crate::types::interface::Config; +use crate::types::state::CollectionType; +use crate::usage::types::interface::SetUserUsage; +use crate::usage::types::state::UserUsage; use crate::version::SATELLITE_VERSION; use ic_cdk::api::trap; use ic_cdk_macros::{init, post_upgrade, pre_upgrade, query, update}; @@ -33,7 +37,7 @@ use junobuild_shared::types::interface::{ }; use junobuild_shared::types::list::ListParams; use junobuild_shared::types::list::ListResults; -use junobuild_shared::types::state::Controllers; +use junobuild_shared::types::state::{Controllers, UserId}; use junobuild_storage::http::types::{ HttpRequest, HttpResponse, StreamingCallbackHttpResponse, StreamingCallbackToken, }; @@ -393,6 +397,31 @@ pub fn get_many_assets( satellite::get_many_assets(assets) } +// --------------------------------------------------------- +// User usage +// --------------------------------------------------------- + +#[doc(hidden)] +#[query] +pub fn get_user_usage( + collection_key: CollectionKey, + collection_type: CollectionType, + user_id: Option, +) -> Option { + satellite::get_user_usage(&collection_key, &collection_type, &user_id) +} + +#[doc(hidden)] +#[update(guard = "caller_is_admin_controller")] +pub fn set_user_usage( + collection_key: CollectionKey, + collection_type: CollectionType, + user_id: UserId, + usage: SetUserUsage, +) -> UserUsage { + satellite::set_user_usage(&collection_key, &collection_type, &user_id, &usage) +} + // --------------------------------------------------------- // Mgmt // --------------------------------------------------------- @@ -440,11 +469,11 @@ macro_rules! include_satellite { count_docs, del_asset, del_assets, del_controllers, del_custom_domain, del_doc, del_docs, del_filtered_assets, del_filtered_docs, del_many_assets, del_many_docs, del_rule, deposit_cycles, get_asset, get_auth_config, get_config, get_db_config, - get_doc, get_many_assets, get_many_docs, get_storage_config, http_request, - http_request_streaming_callback, init, init_asset_upload, list_assets, + get_doc, get_many_assets, get_many_docs, get_storage_config, get_user_usage, + http_request, http_request_streaming_callback, init, init_asset_upload, list_assets, list_controllers, list_custom_domains, list_docs, list_rules, memory_size, post_upgrade, pre_upgrade, set_auth_config, set_controllers, set_custom_domain, - set_db_config, set_doc, set_many_docs, set_rule, set_storage_config, + set_db_config, set_doc, set_many_docs, set_rule, set_storage_config, set_user_usage, upload_asset_chunk, version, }; diff --git a/src/libs/satellite/src/memory.rs b/src/libs/satellite/src/memory.rs index 78e1acb99..fb751d7e1 100644 --- a/src/libs/satellite/src/memory.rs +++ b/src/libs/satellite/src/memory.rs @@ -9,6 +9,7 @@ const UPGRADES: MemoryId = MemoryId::new(0); const DB: MemoryId = MemoryId::new(1); const ASSETS: MemoryId = MemoryId::new(2); const CONTENT_CHUNKS: MemoryId = MemoryId::new(3); +const USER_USAGE: MemoryId = MemoryId::new(4); thread_local! { pub static STATE: RefCell = RefCell::default(); @@ -33,10 +34,15 @@ fn get_memory_content_chunks() -> Memory { MEMORY_MANAGER.with(|m| m.borrow().get(CONTENT_CHUNKS)) } +fn get_memory_user_usage() -> Memory { + MEMORY_MANAGER.with(|m| m.borrow().get(USER_USAGE)) +} + pub fn init_stable_state() -> StableState { StableState { db: StableBTreeMap::init(get_memory_db()), assets: StableBTreeMap::init(get_memory_assets()), content_chunks: StableBTreeMap::init(get_memory_content_chunks()), + user_usage: StableBTreeMap::init(get_memory_user_usage()), } } diff --git a/src/libs/satellite/src/satellite.rs b/src/libs/satellite/src/satellite.rs index 54231f8ca..611d65c15 100644 --- a/src/libs/satellite/src/satellite.rs +++ b/src/libs/satellite/src/satellite.rs @@ -36,8 +36,15 @@ use crate::storage::store::{ set_domain_store, }; use crate::storage::strategy_impls::StorageState; -use crate::types::interface::{Config, CollectionType}; -use crate::types::state::{HeapState, RuntimeState, State}; +use crate::types::interface::Config; +use crate::types::state::{CollectionType, HeapState, RuntimeState, State}; +use crate::usage::types::interface::SetUserUsage; +use crate::usage::types::state::UserUsage; +use crate::usage::user_usage::{ + decrease_db_usage, decrease_db_usage_by, decrease_storage_usage, decrease_storage_usage_by, + get_db_usage_by_id, get_storage_usage_by_id, increase_db_usage, increase_storage_usage, + set_db_usage, set_storage_usage, +}; use ciborium::{from_reader, into_writer}; use ic_cdk::api::call::{arg_data, ArgDecoderConfig}; use ic_cdk::api::{caller, trap}; @@ -54,7 +61,7 @@ use junobuild_shared::types::interface::{DeleteControllersArgs, SegmentArgs, Set use junobuild_shared::types::list::ListParams; use junobuild_shared::types::list::ListResults; use junobuild_shared::types::memory::Memory; -use junobuild_shared::types::state::{ControllerScope, Controllers}; +use junobuild_shared::types::state::{ControllerScope, Controllers, UserId}; use junobuild_shared::upgrade::{read_post_upgrade, write_pre_upgrade}; use junobuild_storage::http::types::{ HttpRequest, HttpResponse, StreamingCallbackHttpResponse, StreamingCallbackToken, @@ -120,10 +127,12 @@ pub fn post_upgrade() { pub fn set_doc(collection: CollectionKey, key: Key, doc: SetDoc) -> Doc { let caller = caller(); - let result = set_doc_store(caller, collection, key, doc); + let result = set_doc_store(caller, collection.clone(), key, doc); match result { Ok(doc) => { + increase_db_usage(&collection, &caller); + invoke_on_set_doc(&caller, &doc); doc.data.after @@ -146,7 +155,10 @@ pub fn get_doc(collection: CollectionKey, key: Key) -> Option { pub fn del_doc(collection: CollectionKey, key: Key, doc: DelDoc) { let caller = caller(); - let deleted_doc = delete_doc_store(caller, collection, key, doc).unwrap_or_else(|e| trap(&e)); + let deleted_doc = + delete_doc_store(caller, collection.clone(), key, doc).unwrap_or_else(|e| trap(&e)); + + decrease_db_usage(&collection, &caller); invoke_on_delete_doc(&caller, &deleted_doc); } @@ -189,8 +201,10 @@ pub fn set_many_docs(docs: Vec<(CollectionKey, Key, SetDoc)>) -> Vec<(Key, Doc)> let mut results: Vec<(Key, Doc)> = Vec::new(); for (collection, key, doc) in docs { - let result = - set_doc_store(caller, collection, key.clone(), doc).unwrap_or_else(|e| trap(&e)); + let result = set_doc_store(caller, collection.clone(), key.clone(), doc) + .unwrap_or_else(|e| trap(&e)); + + increase_db_usage(&collection, &caller); results.push((result.key.clone(), result.data.after.clone())); @@ -208,8 +222,11 @@ pub fn del_many_docs(docs: Vec<(CollectionKey, Key, DelDoc)>) { let mut results: Vec>> = Vec::new(); for (collection, key, doc) in docs { - let deleted_doc = - delete_doc_store(caller, collection, key.clone(), doc).unwrap_or_else(|e| trap(&e)); + let deleted_doc = delete_doc_store(caller, collection.clone(), key.clone(), doc) + .unwrap_or_else(|e| trap(&e)); + + decrease_db_usage(&collection, &caller); + results.push(deleted_doc); } @@ -219,8 +236,10 @@ pub fn del_many_docs(docs: Vec<(CollectionKey, Key, DelDoc)>) { pub fn del_filtered_docs(collection: CollectionKey, filter: ListParams) { let caller = caller(); - let results = - delete_filtered_docs_store(caller, collection, &filter).unwrap_or_else(|e| trap(&e)); + let results = delete_filtered_docs_store(caller, collection.clone(), &filter) + .unwrap_or_else(|e| trap(&e)); + + decrease_db_usage_by(&collection, &caller, results.len() as u32); invoke_on_delete_filtered_docs(&caller, &results); } @@ -429,6 +448,8 @@ pub fn commit_asset_upload(commit: CommitBatch) { let asset = commit_batch_store(caller, commit).unwrap_or_else(|e| trap(&e)); + increase_storage_usage(&asset.key.collection, &caller); + invoke_upload_asset(&caller, &asset); } @@ -460,7 +481,11 @@ pub fn del_asset(collection: CollectionKey, full_path: FullPath) { let result = delete_asset_store(caller, &collection, full_path); match result { - Ok(asset) => invoke_on_delete_asset(&caller, &asset), + Ok(asset) => { + decrease_storage_usage(&collection, &caller); + + invoke_on_delete_asset(&caller, &asset) + } Err(error) => trap(&["Asset cannot be deleted: ", &error].join("")), } } @@ -473,6 +498,9 @@ pub fn del_many_assets(assets: Vec<(CollectionKey, String)>) { for (collection, full_path) in assets { let deleted_asset = delete_asset_store(caller, &collection, full_path).unwrap_or_else(|e| trap(&e)); + + decrease_storage_usage(&collection, &caller); + results.push(deleted_asset); } @@ -482,8 +510,10 @@ pub fn del_many_assets(assets: Vec<(CollectionKey, String)>) { pub fn del_filtered_assets(collection: CollectionKey, filter: ListParams) { let caller = caller(); - let results = - delete_filtered_assets_store(caller, collection, &filter).unwrap_or_else(|e| trap(&e)); + let results = delete_filtered_assets_store(caller, collection.clone(), &filter) + .unwrap_or_else(|e| trap(&e)); + + decrease_storage_usage_by(&collection, &caller, results.len() as u32); invoke_on_delete_filtered_assets(&caller, &results); } @@ -528,3 +558,37 @@ pub fn get_many_assets( }) .collect() } + +// --------------------------------------------------------- +// User usage +// --------------------------------------------------------- + +pub fn get_user_usage( + collection: &CollectionKey, + collection_type: &CollectionType, + user_id: &Option, +) -> Option { + let caller = caller(); + let user_id_or_caller = user_id.unwrap_or(caller); + + match collection_type { + CollectionType::Db => get_db_usage_by_id(caller, collection, &user_id_or_caller), + CollectionType::Storage => get_storage_usage_by_id(caller, collection, &user_id_or_caller), + } +} + +pub fn set_user_usage( + collection: &CollectionKey, + collection_type: &CollectionType, + user_id: &UserId, + usage: &SetUserUsage, +) -> UserUsage { + match collection_type { + CollectionType::Db => { + set_db_usage(collection, user_id, usage.items_count).unwrap_or_else(|e| trap(&e)) + } + CollectionType::Storage => { + set_storage_usage(collection, user_id, usage.items_count).unwrap_or_else(|e| trap(&e)) + } + } +} diff --git a/src/libs/satellite/src/types.rs b/src/libs/satellite/src/types.rs index 91d24ae5d..e295a5981 100644 --- a/src/libs/satellite/src/types.rs +++ b/src/libs/satellite/src/types.rs @@ -3,6 +3,7 @@ pub mod state { use crate::db::types::state::{DbHeapState, DbRuntimeState, DbStable}; use crate::memory::init_stable_state; use crate::storage::types::state::{AssetsStable, ContentChunksStable}; + use crate::usage::types::state::UserUsageStable; use candid::CandidType; use junobuild_shared::types::state::Controllers; use junobuild_storage::types::state::StorageHeapState; @@ -27,6 +28,7 @@ pub mod state { pub db: DbStable, pub assets: AssetsStable, pub content_chunks: ContentChunksStable, + pub user_usage: UserUsageStable, } #[derive(Default, CandidType, Serialize, Deserialize, Clone)] @@ -42,6 +44,12 @@ pub mod state { pub rng: Option, // rng = Random Number Generator pub db: DbRuntimeState, } + + #[derive(CandidType, Deserialize, Serialize, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)] + pub enum CollectionType { + Db, + Storage, + } } pub mod interface { @@ -51,12 +59,6 @@ pub mod interface { use junobuild_storage::types::config::StorageConfig; use serde::Deserialize; - #[derive(CandidType, Deserialize)] - pub enum CollectionType { - Db, - Storage, - } - #[derive(CandidType, Deserialize)] pub struct Config { pub storage: StorageConfig, diff --git a/src/libs/satellite/src/usage/impls.rs b/src/libs/satellite/src/usage/impls.rs new file mode 100644 index 000000000..bee01a7db --- /dev/null +++ b/src/libs/satellite/src/usage/impls.rs @@ -0,0 +1,100 @@ +use crate::types::state::CollectionType; +use crate::usage::types::interface::ModificationType; +use crate::usage::types::state::{UserUsage, UserUsageKey}; +use ic_cdk::api::time; +use ic_stable_structures::storable::Bound; +use ic_stable_structures::Storable; +use junobuild_collections::types::core::CollectionKey; +use junobuild_shared::constants::INITIAL_VERSION; +use junobuild_shared::serializers::{deserialize_from_bytes, serialize_to_bytes}; +use junobuild_shared::types::state::{Timestamp, UserId, Version}; +use std::borrow::Cow; + +impl Storable for UserUsage { + fn to_bytes(&self) -> Cow<[u8]> { + serialize_to_bytes(self) + } + + fn from_bytes(bytes: Cow<[u8]>) -> Self { + deserialize_from_bytes(bytes) + } + + const BOUND: Bound = Bound::Unbounded; +} + +impl Storable for UserUsageKey { + fn to_bytes(&self) -> Cow<[u8]> { + serialize_to_bytes(self) + } + + fn from_bytes(bytes: Cow<[u8]>) -> Self { + deserialize_from_bytes(bytes) + } + + const BOUND: Bound = Bound::Unbounded; +} + +impl UserUsage { + pub fn set(current_user_usage: &Option, count: u32) -> Self { + UserUsage::apply_update(current_user_usage, count) + } + + pub fn increase_or_decrease( + current_user_usage: &Option, + modification: &ModificationType, + count: Option, + ) -> Self { + let count = count.unwrap_or(1); + + // User usage for the collection + + let items_count: u32 = match current_user_usage { + None => 1, + Some(current_user_usage) => match modification { + ModificationType::Set => current_user_usage.items_count.saturating_add(count), + ModificationType::Delete => current_user_usage.items_count.saturating_sub(count), + }, + }; + + UserUsage::apply_update(current_user_usage, items_count) + } + + fn apply_update(current_user_usage: &Option, items_count: u32) -> Self { + let now = time(); + + // Metadata for the UserUsage entity entry + + let created_at: Timestamp = match current_user_usage { + None => now, + Some(current_user_usage) => current_user_usage.created_at, + }; + + let version: Version = match current_user_usage { + None => INITIAL_VERSION, + Some(current_user_usage) => current_user_usage.version.unwrap_or_default() + 1, + }; + + let updated_at: Timestamp = now; + + UserUsage { + items_count, + created_at, + updated_at, + version: Some(version), + } + } +} + +impl UserUsageKey { + pub fn create( + user_id: &UserId, + collection_key: &CollectionKey, + collection_type: &CollectionType, + ) -> Self { + Self { + user_id: *user_id, + collection_key: collection_key.clone(), + collection_type: collection_type.clone(), + } + } +} diff --git a/src/libs/satellite/src/usage/mod.rs b/src/libs/satellite/src/usage/mod.rs new file mode 100644 index 000000000..08773a8d9 --- /dev/null +++ b/src/libs/satellite/src/usage/mod.rs @@ -0,0 +1,4 @@ +mod impls; +mod store; +pub mod types; +pub mod user_usage; diff --git a/src/libs/satellite/src/usage/store.rs b/src/libs/satellite/src/usage/store.rs new file mode 100644 index 000000000..22e387202 --- /dev/null +++ b/src/libs/satellite/src/usage/store.rs @@ -0,0 +1,103 @@ +use crate::memory::STATE; +use crate::types::state::CollectionType; +use crate::usage::types::interface::ModificationType; +use crate::usage::types::state::{UserUsage, UserUsageKey, UserUsageStable}; +use junobuild_collections::types::core::CollectionKey; +use junobuild_shared::types::state::UserId; + +pub fn get_user_usage( + collection_key: &CollectionKey, + collection_type: &CollectionType, + user_id: &UserId, +) -> Option { + STATE.with(|state| { + get_user_usage_impl( + collection_key, + collection_type, + user_id, + &state.borrow().stable.user_usage, + ) + }) +} + +pub fn update_user_usage( + collection_key: &CollectionKey, + collection_type: &CollectionType, + user_id: &UserId, + modification: &ModificationType, + count: Option, +) { + STATE.with(|state| { + update_user_usage_impl( + collection_key, + collection_type, + user_id, + modification, + count, + &mut state.borrow_mut().stable.user_usage, + ) + }) +} + +pub fn set_user_usage( + collection_key: &CollectionKey, + collection_type: &CollectionType, + user_id: &UserId, + count: u32, +) -> UserUsage { + STATE.with(|state| { + set_user_usage_impl( + collection_key, + collection_type, + user_id, + count, + &mut state.borrow_mut().stable.user_usage, + ) + }) +} + +fn get_user_usage_impl( + collection_key: &CollectionKey, + collection_type: &CollectionType, + user_id: &UserId, + state: &UserUsageStable, +) -> Option { + let key = UserUsageKey::create(user_id, collection_key, collection_type); + + state.get(&key) +} + +fn update_user_usage_impl( + collection_key: &CollectionKey, + collection_type: &CollectionType, + user_id: &UserId, + modification: &ModificationType, + count: Option, + state: &mut UserUsageStable, +) { + let key = UserUsageKey::create(user_id, collection_key, collection_type); + + let current_usage = state.get(&key); + + let update_usage = UserUsage::increase_or_decrease(¤t_usage, modification, count); + + state.insert(key, update_usage); +} + +fn set_user_usage_impl( + collection_key: &CollectionKey, + collection_type: &CollectionType, + user_id: &UserId, + count: u32, + state: &mut UserUsageStable, +) -> UserUsage { + let key = UserUsageKey::create(user_id, collection_key, collection_type); + + let current_usage = state.get(&key); + + let update_usage = UserUsage::set(¤t_usage, count); + + state.insert(key, update_usage.clone()); + + update_usage +} diff --git a/src/libs/satellite/src/usage/types.rs b/src/libs/satellite/src/usage/types.rs new file mode 100644 index 000000000..6044e9332 --- /dev/null +++ b/src/libs/satellite/src/usage/types.rs @@ -0,0 +1,61 @@ +pub mod state { + use crate::types::state::CollectionType; + use candid::{CandidType, Deserialize}; + use ic_stable_structures::StableBTreeMap; + use junobuild_collections::types::core::CollectionKey; + use junobuild_shared::types::memory::Memory; + use junobuild_shared::types::state::{Timestamp, UserId, Version}; + use serde::Serialize; + + pub type UserUsageStable = StableBTreeMap; + + /// A unique key for identifying user usage within a collection. + /// + /// It consists of: + /// - `user_id`: The unique identifier for the user which is matched to the caller. + /// - `collection_key`: The collection where usage is tracked. + /// - `collection_type`: The type of collection (`Db` for datastore, `Storage` for assets). + #[derive(CandidType, Serialize, Deserialize, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)] + pub struct UserUsageKey { + pub user_id: UserId, + pub collection_key: CollectionKey, + pub collection_type: CollectionType, + } + + /// Tracks the usage (create, set and delete) of a user in a collection. + /// + /// + /// Fields: + /// - `items_count`: The total number of changes by the user. Is not necessary equals to the number of items due to the lack of migration and the fact that controllers can reset the value. + /// - `created_at`: The timestamp when this user usage entry was first recorded. + /// - `updated_at`: The timestamp of the last update to this user usage entry. + /// - `version`: An optional field representing the version of this usage entry. In the future we might implement checks to avoid overwrite but, this is not the case currently. + #[derive(CandidType, Serialize, Deserialize, Clone)] + pub struct UserUsage { + pub items_count: u32, + pub created_at: Timestamp, + pub updated_at: Timestamp, + pub version: Option, + } +} + +pub mod interface { + use candid::{CandidType, Deserialize}; + use serde::Serialize; + + pub enum ModificationType { + Set, + Delete, + } + + /// Represents the parameters for setting or updating a user's usage entry for a controller. + /// + /// This is useful if one want to set a value after the upgrade, given the lack of migration, or if a controller ever wants to reset the value to allow a user who would hit the limit to continue submitted changes. + /// + /// It includes: + /// - `items_count`: The total number of changes the user has in a specific collection. + #[derive(Default, CandidType, Serialize, Deserialize, Clone)] + pub struct SetUserUsage { + pub items_count: u32, + } +} diff --git a/src/libs/satellite/src/usage/user_usage.rs b/src/libs/satellite/src/usage/user_usage.rs new file mode 100644 index 000000000..37c23bf15 --- /dev/null +++ b/src/libs/satellite/src/usage/user_usage.rs @@ -0,0 +1,172 @@ +use crate::get_controllers; +use crate::types::state::CollectionType; +use crate::usage::store::{ + get_user_usage as get_user_usage_store, set_user_usage, update_user_usage, +}; +use crate::usage::types::interface::ModificationType; +use crate::usage::types::state::UserUsage; +use candid::Principal; +use junobuild_collections::constants::{ + ASSETS_COLLECTIONS_NO_USER_USAGE, DB_COLLECTIONS_NO_USER_USAGE, +}; +use junobuild_collections::types::core::CollectionKey; +use junobuild_shared::controllers::is_controller; +use junobuild_shared::types::state::{Controllers, UserId}; +use junobuild_shared::utils::principal_not_anonymous_and_equal; + +pub fn get_db_usage_by_id( + caller: Principal, + collection: &CollectionKey, + user_id: &UserId, +) -> Option { + get_user_usage_by_id(caller, collection, &CollectionType::Db, user_id) +} + +pub fn get_storage_usage_by_id( + caller: Principal, + collection: &CollectionKey, + user_id: &UserId, +) -> Option { + get_user_usage_by_id(caller, collection, &CollectionType::Storage, user_id) +} + +fn get_user_usage_by_id( + caller: Principal, + collection_key: &CollectionKey, + collection_type: &CollectionType, + user_id: &UserId, +) -> Option { + let controllers: Controllers = get_controllers(); + + if principal_not_anonymous_and_equal(*user_id, caller) || is_controller(caller, &controllers) { + return get_user_usage_store(collection_key, collection_type, user_id); + } + + None +} + +pub fn increase_db_usage(collection: &CollectionKey, user_id: &UserId) { + if is_db_collection_no_usage(collection) { + return; + } + + update_user_usage( + collection, + &CollectionType::Db, + user_id, + &ModificationType::Set, + None, + ); +} + +pub fn set_db_usage( + collection: &CollectionKey, + user_id: &UserId, + count: u32, +) -> Result { + if is_db_collection_no_usage(collection) { + return Err(format!( + "Datastore usage is not recorded for collection {}.", + collection + )); + } + + let usage = set_user_usage(collection, &CollectionType::Db, user_id, count); + + Ok(usage) +} + +pub fn decrease_db_usage(collection: &CollectionKey, user_id: &UserId) { + if is_db_collection_no_usage(collection) { + return; + } + + update_user_usage( + collection, + &CollectionType::Db, + user_id, + &ModificationType::Delete, + None, + ); +} + +pub fn decrease_db_usage_by(collection: &CollectionKey, user_id: &UserId, count: u32) { + if is_db_collection_no_usage(collection) { + return; + } + + update_user_usage( + collection, + &CollectionType::Db, + user_id, + &ModificationType::Delete, + Some(count), + ); +} + +pub fn increase_storage_usage(collection: &CollectionKey, user_id: &UserId) { + if is_storage_collection_no_usage(collection) { + return; + } + + update_user_usage( + collection, + &CollectionType::Storage, + user_id, + &ModificationType::Set, + None, + ); +} + +pub fn set_storage_usage( + collection: &CollectionKey, + user_id: &UserId, + count: u32, +) -> Result { + if is_storage_collection_no_usage(collection) { + return Err(format!( + "Storage usage is not recorded for collection {}.", + collection + )); + } + + let usage = set_user_usage(collection, &CollectionType::Storage, user_id, count); + + Ok(usage) +} + +pub fn decrease_storage_usage(collection: &CollectionKey, user_id: &UserId) { + if is_storage_collection_no_usage(collection) { + return; + } + + update_user_usage( + collection, + &CollectionType::Storage, + user_id, + &ModificationType::Delete, + None, + ); +} + +pub fn decrease_storage_usage_by(collection: &CollectionKey, user_id: &UserId, count: u32) { + if is_storage_collection_no_usage(collection) { + return; + } + + update_user_usage( + collection, + &CollectionType::Storage, + user_id, + &ModificationType::Delete, + Some(count), + ); +} + +fn is_db_collection_no_usage(collection: &CollectionKey) -> bool { + DB_COLLECTIONS_NO_USER_USAGE.contains(&collection.as_str()) +} + +fn is_storage_collection_no_usage(collection: &CollectionKey) -> bool { + ASSETS_COLLECTIONS_NO_USER_USAGE.contains(&collection.as_str()) +} diff --git a/src/satellite/satellite.did b/src/satellite/satellite.did index a275796e7..d36243651 100644 --- a/src/satellite/satellite.did +++ b/src/satellite/satellite.did @@ -157,6 +157,7 @@ type SetRule = record { rate_config : opt RateConfig; write : Permission; }; +type SetUserUsage = record { items_count : nat32 }; type StorageConfig = record { iframe : opt StorageConfigIFrame; rewrites : vec record { text; text }; @@ -199,6 +200,12 @@ type UploadChunk = record { order_id : opt nat; }; type UploadChunkResult = record { chunk_id : nat }; +type UserUsage = record { + updated_at : nat64; + created_at : nat64; + version : opt nat64; + items_count : nat32; +}; service : () -> { commit_asset_upload : (CommitBatch) -> (); count_assets : (text, ListParams) -> (nat64) query; @@ -232,6 +239,9 @@ service : () -> { ) query; get_rule : (CollectionType, text) -> (opt Rule) query; get_storage_config : () -> (StorageConfig) query; + get_user_usage : (text, CollectionType, opt principal) -> ( + opt UserUsage, + ) query; http_request : (HttpRequest) -> (HttpResponse) query; http_request_streaming_callback : (StreamingCallbackToken) -> ( StreamingCallbackHttpResponse, @@ -255,6 +265,9 @@ service : () -> { ); set_rule : (CollectionType, text, SetRule) -> (Rule); set_storage_config : (StorageConfig) -> (); + set_user_usage : (text, CollectionType, principal, SetUserUsage) -> ( + UserUsage, + ); upload_asset_chunk : (UploadChunk) -> (UploadChunkResult); version : () -> (text) query; } diff --git a/src/tests/satellite.datastore.spec.ts b/src/tests/satellite.datastore.spec.ts index 634771a81..65a9a2575 100644 --- a/src/tests/satellite.datastore.spec.ts +++ b/src/tests/satellite.datastore.spec.ts @@ -767,7 +767,7 @@ describe.each([{ memory: { Heap: null } }, { memory: { Stable: null } }])( }, { memory: { Stable: null }, - expectMemory: 25_231_360n + expectMemory: 33_619_968n } ])('With collection', ({ memory, expectMemory }) => { const errorMsg = `${'Heap' in memory ? 'Heap' : 'Stable'} memory usage exceeded: ${expectMemory} bytes used, 20000 bytes allowed.`; diff --git a/src/tests/satellite.storage.spec.ts b/src/tests/satellite.storage.spec.ts index 971ee196e..ff76146f1 100644 --- a/src/tests/satellite.storage.spec.ts +++ b/src/tests/satellite.storage.spec.ts @@ -1165,7 +1165,7 @@ describe('Satellite storage', () => { }, { memory: { Stable: null }, - expectMemory: 25_231_360n, + expectMemory: 33_619_968n, allowedMemory: maxStableMemorySize, preUploadCount: 0 } diff --git a/src/tests/satellite.user-usage.spec.ts b/src/tests/satellite.user-usage.spec.ts new file mode 100644 index 000000000..00e56a9d7 --- /dev/null +++ b/src/tests/satellite.user-usage.spec.ts @@ -0,0 +1,596 @@ +import type { + DelDoc, + ListParams, + _SERVICE as SatelliteActor, + SetDoc, + SetRule, + UserUsage +} from '$declarations/satellite/satellite.did'; +import { idlFactory as idlFactorSatellite } from '$declarations/satellite/satellite.factory.did'; +import { Ed25519KeyIdentity } from '@dfinity/identity'; +import type { Principal } from '@dfinity/principal'; +import { assertNonNullish, fromNullable, toNullable } from '@dfinity/utils'; +import { type Actor, PocketIc } from '@hadronous/pic'; +import { toArray } from '@junobuild/utils'; +import { nanoid } from 'nanoid'; +import { beforeAll, describe, expect, inject } from 'vitest'; +import { SATELLITE_ADMIN_ERROR_MSG } from './constants/satellite-tests.constants'; +import { uploadAsset } from './utils/satellite-storage-tests.utils'; +import { controllersInitArgs, SATELLITE_WASM_PATH } from './utils/setup-tests.utils'; + +describe('Satellite User Usage', () => { + let pic: PocketIc; + let actor: Actor; + + const controller = Ed25519KeyIdentity.generate(); + + const TEST_COLLECTION = 'test'; + + const setRule: SetRule = { + memory: toNullable({ Heap: null }), + max_size: toNullable(), + max_capacity: toNullable(), + read: { Managed: null }, + mutable_permissions: toNullable(), + write: { Managed: null }, + version: toNullable(), + rate_config: toNullable() + }; + + const NO_FILTER_PARAMS: ListParams = { + matcher: toNullable(), + order: toNullable(), + owner: toNullable(), + paginate: toNullable() + }; + + beforeAll(async () => { + pic = await PocketIc.create(inject('PIC_URL')); + + const { actor: c } = await pic.setupCanister({ + idlFactory: idlFactorSatellite, + wasm: SATELLITE_WASM_PATH, + arg: controllersInitArgs(controller), + sender: controller.getPrincipal() + }); + + actor = c; + + actor.setIdentity(controller); + }); + + afterAll(async () => { + await pic?.tearDown(); + }); + + describe('Datastore', async () => { + const data = await toArray({ + hello: 'World' + }); + + const COLLECTION_TYPE = { Db: null }; + + beforeAll(async () => { + const { set_rule } = actor; + await set_rule(COLLECTION_TYPE, TEST_COLLECTION, setRule); + }); + + const createDoc = async (): Promise => { + const key = nanoid(); + + const { set_doc } = actor; + + await set_doc(TEST_COLLECTION, key, { + data, + description: toNullable(), + version: toNullable() + }); + + return key; + }; + + const user = Ed25519KeyIdentity.generate(); + let countTotalTestVersion: number; + + describe('User', () => { + beforeAll(() => { + actor.setIdentity(user); + }); + + const countSetDocs = 10; + + it('should get a usage count after set documents', async () => { + await Promise.all(Array.from({ length: countSetDocs }).map(createDoc)); + + const { get_user_usage } = actor; + + const usageResponse = await get_user_usage(TEST_COLLECTION, COLLECTION_TYPE, toNullable()); + + const usage = fromNullable(usageResponse); + + assertNonNullish(usage); + + expect(usage.items_count).toEqual(countSetDocs); + + expect(usage.updated_at).not.toBeUndefined(); + expect(usage.updated_at).toBeGreaterThan(0n); + expect(usage.created_at).not.toBeUndefined(); + expect(usage.created_at).toBeGreaterThan(0n); + expect(usage.updated_at).toBeGreaterThan(usage.created_at); + + expect(usage.version).toEqual(toNullable(BigInt(countSetDocs))); + }); + + const countSetManyDocs = 5; + + it('should get a usage count after set many documents', async () => { + const { set_many_docs } = actor; + + const docs: [string, string, SetDoc][] = Array.from({ length: countSetManyDocs }).map( + () => [ + TEST_COLLECTION, + nanoid(), + { + data, + description: toNullable(), + version: toNullable() + } + ] + ); + + await set_many_docs(docs); + + const { get_user_usage } = actor; + + const usageResponse = await get_user_usage(TEST_COLLECTION, COLLECTION_TYPE, toNullable()); + + const usage = fromNullable(usageResponse); + + assertNonNullish(usage); + + expect(usage.items_count).toEqual(countSetManyDocs + countSetDocs); + expect(usage.version).toEqual(toNullable(BigInt(countSetManyDocs + countSetDocs))); + }); + + const countDelDoc = 1; + + it('should get a usage count after delete document', async () => { + const { del_doc, list_docs } = actor; + + const { items } = await list_docs(TEST_COLLECTION, NO_FILTER_PARAMS); + + const doc = items[0][1]; + + await del_doc(TEST_COLLECTION, items[0][0], { + version: doc.version ?? [] + }); + + const { get_user_usage } = actor; + + const usageResponse = await get_user_usage(TEST_COLLECTION, COLLECTION_TYPE, toNullable()); + + const usage = fromNullable(usageResponse); + + assertNonNullish(usage); + + expect(usage.items_count).toEqual(countSetManyDocs + countSetDocs - countDelDoc); + expect(usage.version).toEqual( + toNullable(BigInt(countSetManyDocs + countSetDocs + countDelDoc)) + ); + }); + + const countDelManyDocs = 2; + + it('should get a usage count after delete many documents', async () => { + const { del_many_docs, list_docs } = actor; + + const { items } = await list_docs(TEST_COLLECTION, NO_FILTER_PARAMS); + + const docs: [string, string, DelDoc][] = [items[0], items[1]].map(([key, doc]) => [ + TEST_COLLECTION, + key, + { + version: doc.version ?? [] + } + ]); + + await del_many_docs(docs); + + const { get_user_usage } = actor; + + const usageResponse = await get_user_usage(TEST_COLLECTION, COLLECTION_TYPE, toNullable()); + + const usage = fromNullable(usageResponse); + + assertNonNullish(usage); + + expect(usage.items_count).toEqual( + countSetManyDocs + countSetDocs - countDelDoc - countDelManyDocs + ); + expect(usage.version).toEqual( + toNullable(BigInt(countSetManyDocs + countSetDocs + countDelDoc + countDelManyDocs)) + ); + }); + + it('should get a usage count after delete filtered docs', async () => { + const { del_filtered_docs } = actor; + + await del_filtered_docs(TEST_COLLECTION, NO_FILTER_PARAMS); + + const { get_user_usage } = actor; + + const usageResponse = await get_user_usage(TEST_COLLECTION, COLLECTION_TYPE, toNullable()); + + const usage = fromNullable(usageResponse); + + assertNonNullish(usage); + + countTotalTestVersion = + countSetManyDocs + countSetDocs + countDelDoc + countDelManyDocs + 1; + + expect(usage.items_count).toEqual(0); + expect(usage.version).toEqual(toNullable(BigInt(countTotalTestVersion))); + }); + }); + + describe('Guards', () => { + const user1 = Ed25519KeyIdentity.generate(); + + const fetchUsage = async (userId?: Principal): Promise => { + const { get_user_usage } = actor; + + const usageResponse = await get_user_usage( + TEST_COLLECTION, + COLLECTION_TYPE, + toNullable(userId) + ); + + return fromNullable(usageResponse); + }; + + it('should not get usage of another user', async () => { + actor.setIdentity(user1); + + await createDoc(); + + const usage = await fetchUsage(); + + assertNonNullish(usage); + + expect(usage.items_count).toEqual(1); + + const user2 = Ed25519KeyIdentity.generate(); + + actor.setIdentity(user2); + + const usage2 = await fetchUsage(user1.getPrincipal()); + + expect(usage2).toBeUndefined(); + }); + + it('should get usage of user if controller', async () => { + actor.setIdentity(controller); + + const usage = await fetchUsage(user1.getPrincipal()); + + assertNonNullish(usage); + + expect(usage.items_count).toEqual(1); + }); + + it('should throw errors on set usage', async () => { + actor.setIdentity(user1); + + const { set_user_usage } = actor; + + await expect( + set_user_usage(TEST_COLLECTION, COLLECTION_TYPE, user1.getPrincipal(), { + items_count: 345 + }) + ).rejects.toThrow(SATELLITE_ADMIN_ERROR_MSG); + }); + }); + + describe('No user usage', () => { + beforeAll(() => { + actor.setIdentity(controller); + }); + + it('should get no usage of collection is log', async () => { + const { set_doc, get_doc, get_user_usage } = actor; + + const key = nanoid(); + + await set_doc('#log', key, { + data, + description: toNullable(), + version: toNullable() + }); + + const doc = await get_doc('#log', key); + expect(fromNullable(doc)).not.toBeUndefined(); + + const usageResponse = await get_user_usage('#log', COLLECTION_TYPE, toNullable()); + expect(fromNullable(usageResponse)).toBeUndefined(); + }); + }); + + describe('Admin', () => { + beforeAll(() => { + actor.setIdentity(controller); + }); + + it('should set usage for user', async () => { + const { set_user_usage } = actor; + + const usage = await set_user_usage(TEST_COLLECTION, COLLECTION_TYPE, user.getPrincipal(), { + items_count: 345 + }); + + expect(usage.items_count).toEqual(345); + + expect(usage.updated_at).not.toBeUndefined(); + expect(usage.updated_at).toBeGreaterThan(0n); + expect(usage.created_at).not.toBeUndefined(); + expect(usage.created_at).toBeGreaterThan(0n); + expect(usage.updated_at).toBeGreaterThan(usage.created_at); + + expect(usage.version).toEqual(toNullable(BigInt(countTotalTestVersion + 1))); + }); + }); + }); + + describe('Storage', () => { + const COLLECTION_TYPE = { Storage: null }; + + beforeAll(async () => { + actor.setIdentity(controller); + + const { set_rule } = actor; + await set_rule(COLLECTION_TYPE, TEST_COLLECTION, setRule); + }); + + const upload = async (index: number) => { + const name = `hello-${index}.html`; + const full_path = `/${TEST_COLLECTION}/${name}`; + + await uploadAsset({ + full_path, + name, + collection: TEST_COLLECTION, + actor + }); + }; + + const createUser = async (user: Principal) => { + // We need a user entry to upload to the storage, there is a guard + const { set_doc } = actor; + + await set_doc('#user', user.toText(), { + data: await toArray({ + provider: 'internet_identity' + }), + description: toNullable(), + version: toNullable() + }); + }; + + const user = Ed25519KeyIdentity.generate(); + let countTotalTestVersion: number; + + describe('User', () => { + beforeAll(async () => { + actor.setIdentity(user); + + await createUser(user.getPrincipal()); + }); + + const countUploadAssets = 10; + + it('should get a usage count after update asset', async () => { + await Promise.all( + Array.from({ length: countUploadAssets }).map(async (_, i) => await upload(i)) + ); + + const { get_user_usage } = actor; + + const usageResponse = await get_user_usage(TEST_COLLECTION, COLLECTION_TYPE, toNullable()); + + const usage = fromNullable(usageResponse); + + assertNonNullish(usage); + + expect(usage.items_count).toEqual(countUploadAssets); + + expect(usage.updated_at).not.toBeUndefined(); + expect(usage.updated_at).toBeGreaterThan(0n); + expect(usage.created_at).not.toBeUndefined(); + expect(usage.created_at).toBeGreaterThan(0n); + expect(usage.updated_at).toBeGreaterThan(usage.created_at); + + expect(usage.version).toEqual(toNullable(BigInt(countUploadAssets))); + }); + + const countDelAsset = 1; + + it('should get a usage count after delete one asset', async () => { + const { del_asset, list_assets } = actor; + + const { items } = await list_assets(TEST_COLLECTION, NO_FILTER_PARAMS); + + const asset = items[0][1]; + + await del_asset(TEST_COLLECTION, asset.key.full_path); + + const { get_user_usage } = actor; + + const usageResponse = await get_user_usage(TEST_COLLECTION, COLLECTION_TYPE, toNullable()); + + const usage = fromNullable(usageResponse); + + assertNonNullish(usage); + + expect(usage.items_count).toEqual(countUploadAssets - countDelAsset); + expect(usage.version).toEqual(toNullable(BigInt(countUploadAssets + countDelAsset))); + }); + + const countDelManyAssets = 2; + + it('should get a usage count after delete many assets', async () => { + const { del_many_assets, list_assets } = actor; + + const { items } = await list_assets(TEST_COLLECTION, NO_FILTER_PARAMS); + + const assets: [string, string][] = [items[0], items[1]].map(([_, asset]) => [ + TEST_COLLECTION, + asset.key.full_path + ]); + + await del_many_assets(assets); + + const { get_user_usage } = actor; + + const usageResponse = await get_user_usage(TEST_COLLECTION, COLLECTION_TYPE, toNullable()); + + const usage = fromNullable(usageResponse); + + assertNonNullish(usage); + + expect(usage.items_count).toEqual(countUploadAssets - countDelAsset - countDelManyAssets); + expect(usage.version).toEqual( + toNullable(BigInt(countUploadAssets + countDelAsset + countDelManyAssets)) + ); + }); + + it('should get a usage count after delete filtered assets', async () => { + const { del_filtered_assets } = actor; + + await del_filtered_assets(TEST_COLLECTION, NO_FILTER_PARAMS); + + const { get_user_usage } = actor; + + const usageResponse = await get_user_usage(TEST_COLLECTION, COLLECTION_TYPE, toNullable()); + + const usage = fromNullable(usageResponse); + + assertNonNullish(usage); + + countTotalTestVersion = countUploadAssets + countDelAsset + countDelManyAssets + 1; + + expect(usage.items_count).toEqual(0); + expect(usage.version).toEqual(toNullable(BigInt(countTotalTestVersion))); + }); + }); + + describe('Guards', () => { + const user1 = Ed25519KeyIdentity.generate(); + + beforeAll(async () => { + actor.setIdentity(user1); + await createUser(user1.getPrincipal()); + }); + + const fetchUsage = async (userId?: Principal): Promise => { + const { get_user_usage } = actor; + + const usageResponse = await get_user_usage( + TEST_COLLECTION, + COLLECTION_TYPE, + toNullable(userId) + ); + return fromNullable(usageResponse); + }; + + it('should not get usage of another user', async () => { + await upload(100); + + const usage = await fetchUsage(); + + assertNonNullish(usage); + + expect(usage.items_count).toEqual(1); + + const user2 = Ed25519KeyIdentity.generate(); + + actor.setIdentity(user2); + await createUser(user2.getPrincipal()); + + const usage2 = await fetchUsage(user1.getPrincipal()); + + expect(usage2).toBeUndefined(); + }); + + it('should get usage of user if controller', async () => { + actor.setIdentity(controller); + + const usage = await fetchUsage(user1.getPrincipal()); + + assertNonNullish(usage); + + expect(usage.items_count).toEqual(1); + }); + + it('should throw errors on set usage', async () => { + actor.setIdentity(user1); + + const { set_user_usage } = actor; + + await expect( + set_user_usage(TEST_COLLECTION, COLLECTION_TYPE, user1.getPrincipal(), { + items_count: 345 + }) + ).rejects.toThrow(SATELLITE_ADMIN_ERROR_MSG); + }); + }); + + describe('No user usage', () => { + beforeAll(() => { + actor.setIdentity(controller); + }); + + it('should get no usage of collection is dapp', async () => { + const { get_asset, get_user_usage } = actor; + + const name = `index.html`; + const full_path = `/${name}`; + + await uploadAsset({ + full_path, + name, + collection: '#dapp', + actor + }); + + const asset = await get_asset('#dapp', full_path); + expect(fromNullable(asset)).not.toBeUndefined(); + + const usageResponse = await get_user_usage('#dapp', COLLECTION_TYPE, toNullable()); + expect(fromNullable(usageResponse)).toBeUndefined(); + }); + }); + + describe('Admin', () => { + beforeAll(() => { + actor.setIdentity(controller); + }); + + it('should set usage for user', async () => { + const { set_user_usage } = actor; + + const usage = await set_user_usage(TEST_COLLECTION, COLLECTION_TYPE, user.getPrincipal(), { + items_count: 456 + }); + + expect(usage.items_count).toEqual(456); + + expect(usage.updated_at).not.toBeUndefined(); + expect(usage.updated_at).toBeGreaterThan(0n); + expect(usage.created_at).not.toBeUndefined(); + expect(usage.created_at).toBeGreaterThan(0n); + expect(usage.updated_at).toBeGreaterThan(usage.created_at); + + expect(usage.version).toEqual(toNullable(BigInt(countTotalTestVersion + 1))); + }); + }); + }); +});