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 from_bytes(bytes: std::borrow::Cow<'_, [u8]>) -> Self {
47        candid::decode_one(&bytes).unwrap()
48    }
49
50    const BOUND: ic_stable_structures::storable::Bound =
51        ic_stable_structures::storable::Bound::Bounded {
52            max_size: 1024,
53            is_fixed_size: false,
54        };
55}
56
57impl Storable for ToolState {
58    fn to_bytes(&self) -> std::borrow::Cow<'_, [u8]> {
59        std::borrow::Cow::Owned(candid::encode_one(self).unwrap())
60    }
61
62    fn from_bytes(bytes: std::borrow::Cow<'_, [u8]>) -> Self {
63        candid::decode_one(&bytes).unwrap()
64    }
65
66    const BOUND: ic_stable_structures::storable::Bound =
67        ic_stable_structures::storable::Bound::Bounded {
68            max_size: 512,
69            is_fixed_size: false,
70        };
71}
72
73impl Storable for ResourceState {
74    fn to_bytes(&self) -> std::borrow::Cow<'_, [u8]> {
75        std::borrow::Cow::Owned(candid::encode_one(self).unwrap())
76    }
77
78    fn from_bytes(bytes: std::borrow::Cow<'_, [u8]>) -> Self {
79        candid::decode_one(&bytes).unwrap()
80    }
81
82    const BOUND: ic_stable_structures::storable::Bound =
83        ic_stable_structures::storable::Bound::Bounded {
84            max_size: 512,
85            is_fixed_size: false,
86        };
87}
88
89thread_local! {
90    /// Global canister state
91    pub static STATE: RefCell<Option<IcarusCanisterState>> = const { RefCell::new(None) };
92}
93
94impl IcarusCanisterState {
95    pub fn init(config: ServerConfig) {
96        let state = Self {
97            config: StableCell::init(get_memory(MEMORY_ID_CONFIG), config).unwrap(),
98            tools: StableBTreeMap::init(get_memory(MEMORY_ID_TOOLS)),
99            resources: StableBTreeMap::init(get_memory(MEMORY_ID_RESOURCES)),
100        };
101
102        STATE.with(|s| *s.borrow_mut() = Some(state));
103    }
104
105    pub fn with<F, R>(f: F) -> R
106    where
107        F: FnOnce(&IcarusCanisterState) -> R,
108    {
109        STATE.with(|s| {
110            let state = s.borrow();
111            let state_ref = state.as_ref().expect("State not initialized");
112            f(state_ref)
113        })
114    }
115
116    /// Get the canister owner principal
117    pub fn get_owner(&self) -> Principal {
118        self.config.get().owner
119    }
120}
121
122// State should not be cloneable as it contains stable memory structures
123// Use STATE.with() to access the global state instead
124
125/// Access control functions
126/// Assert that the caller is the canister owner
127pub fn assert_owner() {
128    let caller = ic_cdk::caller();
129    IcarusCanisterState::with(|state| {
130        let owner = state.get_owner();
131        if caller != owner {
132            ic_cdk::trap(&format!(
133                "Access denied: caller {} is not the owner {}",
134                caller.to_text(),
135                owner.to_text()
136            ));
137        }
138    });
139}
140
141/// Check if the caller is the canister owner without trapping
142pub fn is_owner(caller: Principal) -> bool {
143    IcarusCanisterState::with(|state| caller == state.get_owner())
144}
145
146/// Get the current owner principal
147pub fn get_owner() -> Principal {
148    IcarusCanisterState::with(|state| state.get_owner())
149}
icarus_canister/state.rs

icarus_canister/
state.rs
