icarus_canister/
state.rs

1//! Canister state management
2
3use crate::memory::{get_memory, MEMORY_ID_CONFIG, MEMORY_ID_RESOURCES, MEMORY_ID_TOOLS};
4use candid::{CandidType, Deserialize, Principal};
5use ic_stable_structures::{StableBTreeMap, StableCell, Storable};
6use serde::Serialize;
7use std::cell::RefCell;
8
9/// Main canister state
10pub struct IcarusCanisterState {
11    pub config: StableCell<ServerConfig, crate::memory::Memory>,
12    pub tools: StableBTreeMap<String, ToolState, crate::memory::Memory>,
13    pub resources: StableBTreeMap<String, ResourceState, crate::memory::Memory>,
14}
15
16/// Server configuration stored in stable memory
17#[derive(Debug, Clone, Serialize, Deserialize, CandidType)]
18pub struct ServerConfig {
19    pub name: String,
20    pub version: String,
21    pub canister_id: Principal,
22    pub owner: Principal,
23}
24
25/// State for individual tools
26#[derive(Debug, Clone, Serialize, Deserialize, CandidType)]
27pub struct ToolState {
28    pub name: String,
29    pub enabled: bool,
30    pub call_count: u64,
31    pub is_query: bool,
32}
33
34/// State for individual resources
35#[derive(Debug, Clone, Serialize, Deserialize, CandidType)]
36pub struct ResourceState {
37    pub uri: String,
38    pub access_count: u64,
39}
40
41impl Storable for ServerConfig {
42    fn to_bytes(&self) -> std::borrow::Cow<'_, [u8]> {
43        std::borrow::Cow::Owned(candid::encode_one(self).unwrap())
44    }
45
46    fn into_bytes(self) -> Vec<u8> {
47        candid::encode_one(&self).unwrap()
48    }
49
50    fn from_bytes(bytes: std::borrow::Cow<'_, [u8]>) -> Self {
51        candid::decode_one(&bytes).unwrap()
52    }
53
54    const BOUND: ic_stable_structures::storable::Bound =
55        ic_stable_structures::storable::Bound::Bounded {
56            max_size: 1024,
57            is_fixed_size: false,
58        };
59}
60
61impl Storable for ToolState {
62    fn to_bytes(&self) -> std::borrow::Cow<'_, [u8]> {
63        std::borrow::Cow::Owned(candid::encode_one(self).unwrap())
64    }
65
66    fn into_bytes(self) -> Vec<u8> {
67        candid::encode_one(&self).unwrap()
68    }
69
70    fn from_bytes(bytes: std::borrow::Cow<'_, [u8]>) -> Self {
71        candid::decode_one(&bytes).unwrap()
72    }
73
74    const BOUND: ic_stable_structures::storable::Bound =
75        ic_stable_structures::storable::Bound::Bounded {
76            max_size: 512,
77            is_fixed_size: false,
78        };
79}
80
81impl Storable for ResourceState {
82    fn to_bytes(&self) -> std::borrow::Cow<'_, [u8]> {
83        std::borrow::Cow::Owned(candid::encode_one(self).unwrap())
84    }
85
86    fn into_bytes(self) -> Vec<u8> {
87        candid::encode_one(&self).unwrap()
88    }
89
90    fn from_bytes(bytes: std::borrow::Cow<'_, [u8]>) -> Self {
91        candid::decode_one(&bytes).unwrap()
92    }
93
94    const BOUND: ic_stable_structures::storable::Bound =
95        ic_stable_structures::storable::Bound::Bounded {
96            max_size: 512,
97            is_fixed_size: false,
98        };
99}
100
101thread_local! {
102    /// Global canister state
103    pub static STATE: RefCell<Option<IcarusCanisterState>> = const { RefCell::new(None) };
104}
105
106impl IcarusCanisterState {
107    pub fn init(config: ServerConfig) {
108        let state = Self {
109            config: StableCell::init(get_memory(MEMORY_ID_CONFIG), config),
110            tools: StableBTreeMap::init(get_memory(MEMORY_ID_TOOLS)),
111            resources: StableBTreeMap::init(get_memory(MEMORY_ID_RESOURCES)),
112        };
113
114        STATE.with(|s| *s.borrow_mut() = Some(state));
115    }
116
117    pub fn with<F, R>(f: F) -> R
118    where
119        F: FnOnce(&IcarusCanisterState) -> R,
120    {
121        STATE.with(|s| {
122            let state = s.borrow();
123            let state_ref = state.as_ref().expect("State not initialized");
124            f(state_ref)
125        })
126    }
127
128    /// Get the canister owner principal
129    pub fn get_owner(&self) -> Principal {
130        self.config.get().owner
131    }
132}
133
134// State should not be cloneable as it contains stable memory structures
135// Use STATE.with() to access the global state instead
136
137/// Access control functions
138/// Assert that the caller is the canister owner
139pub fn assert_owner() {
140    let caller = ic_cdk::api::msg_caller();
141    IcarusCanisterState::with(|state| {
142        let owner = state.get_owner();
143        if caller != owner {
144            ic_cdk::trap(format!(
145                "Access denied: caller {} is not the owner {}",
146                caller.to_text(),
147                owner.to_text()
148            ));
149        }
150    });
151}
152
153/// Check if the caller is the canister owner without trapping
154pub fn is_owner(caller: Principal) -> bool {
155    IcarusCanisterState::with(|state| caller == state.get_owner())
156}
157
158/// Get the current owner principal
159pub fn get_owner() -> Principal {
160    IcarusCanisterState::with(|state| state.get_owner())
161}
icarus_canister/state.rs

icarus_canister/
state.rs
