Skip to main content

canic_core/api/runtime/
install.rs

1use crate::{
2    InternalError, InternalErrorOrigin,
3    cdk::{types::Principal, utils::wasm::get_wasm_hash},
4    dto::error::Error,
5    format::byte_size,
6    ids::CanisterRole,
7    ops::runtime::metrics::wasm_store::{
8        WasmStoreMetricOperation, WasmStoreMetricOutcome, WasmStoreMetricReason,
9        WasmStoreMetricSource, WasmStoreMetrics,
10    },
11};
12use async_trait::async_trait;
13use std::{
14    borrow::Cow,
15    collections::BTreeMap,
16    sync::{Mutex, OnceLock},
17};
18
19///
20/// ApprovedModulePayload
21///
22
23#[derive(Clone, Debug, Eq, PartialEq)]
24pub enum ApprovedModulePayload {
25    Chunked {
26        source_canister: Principal,
27        chunk_hashes: Vec<Vec<u8>>,
28    },
29    Embedded {
30        wasm_module: Cow<'static, [u8]>,
31    },
32}
33
34///
35/// ApprovedModuleSource
36///
37
38#[derive(Clone, Debug, Eq, PartialEq)]
39pub struct ApprovedModuleSource {
40    source_label: String,
41    module_hash: Vec<u8>,
42    payload_size_bytes: u64,
43    payload: ApprovedModulePayload,
44}
45
46impl ApprovedModuleSource {
47    /// Construct one chunk-store-backed module source.
48    #[must_use]
49    pub const fn chunked(
50        source_canister: Principal,
51        source_label: String,
52        module_hash: Vec<u8>,
53        chunk_hashes: Vec<Vec<u8>>,
54        payload_size_bytes: u64,
55    ) -> Self {
56        Self {
57            source_label,
58            module_hash,
59            payload_size_bytes,
60            payload: ApprovedModulePayload::Chunked {
61                source_canister,
62                chunk_hashes,
63            },
64        }
65    }
66
67    /// Construct one embedded module source from already packaged wasm bytes.
68    #[must_use]
69    pub fn embedded(source_label: String, wasm_module: &'static [u8]) -> Self {
70        let payload_size_bytes = wasm_module.len() as u64;
71
72        Self {
73            source_label,
74            module_hash: get_wasm_hash(wasm_module),
75            payload_size_bytes,
76            payload: ApprovedModulePayload::Embedded {
77                wasm_module: Cow::Borrowed(wasm_module),
78            },
79        }
80    }
81
82    /// Return the logical source label used for logs and status output.
83    #[must_use]
84    pub fn source_label(&self) -> &str {
85        &self.source_label
86    }
87
88    /// Return the installable wasm module hash.
89    #[must_use]
90    pub fn module_hash(&self) -> &[u8] {
91        &self.module_hash
92    }
93
94    /// Return the formatted module payload size for logs and status output.
95    #[must_use]
96    pub fn payload_size(&self) -> String {
97        byte_size(self.payload_size_bytes)
98    }
99
100    /// Return the raw payload size in bytes.
101    #[must_use]
102    pub const fn payload_size_bytes(&self) -> u64 {
103        self.payload_size_bytes
104    }
105
106    /// Return the chunk count when the source is chunk-store-backed.
107    #[must_use]
108    pub const fn chunk_count(&self) -> usize {
109        match &self.payload {
110            ApprovedModulePayload::Chunked { chunk_hashes, .. } => chunk_hashes.len(),
111            ApprovedModulePayload::Embedded { .. } => 0,
112        }
113    }
114
115    /// Return the underlying payload representation.
116    #[must_use]
117    pub const fn payload(&self) -> &ApprovedModulePayload {
118        &self.payload
119    }
120}
121
122///
123/// ModuleSourceResolver
124///
125
126#[async_trait]
127pub trait ModuleSourceResolver: Send + Sync {
128    /// Resolve the currently approved install source for one canister role.
129    async fn approved_module_source(
130        &self,
131        role: &CanisterRole,
132    ) -> Result<ApprovedModuleSource, Error>;
133}
134
135static MODULE_SOURCE_RESOLVER: OnceLock<&'static dyn ModuleSourceResolver> = OnceLock::new();
136static EMBEDDED_MODULE_SOURCES: OnceLock<Mutex<BTreeMap<CanisterRole, ApprovedModuleSource>>> =
137    OnceLock::new();
138
139///
140/// ModuleSourceRuntimeApi
141///
142
143pub struct ModuleSourceRuntimeApi;
144
145impl ModuleSourceRuntimeApi {
146    /// Register one built-in module source override for the current process.
147    pub fn register_embedded_module_source(role: CanisterRole, source: ApprovedModuleSource) {
148        let sources = EMBEDDED_MODULE_SOURCES.get_or_init(|| Mutex::new(BTreeMap::new()));
149        let mut sources = sources
150            .lock()
151            .unwrap_or_else(std::sync::PoisonError::into_inner);
152
153        match sources.get(&role) {
154            Some(existing) if existing == &source => {}
155            Some(existing) => {
156                panic!(
157                    "embedded module source for role '{role}' was already registered with a different payload: existing='{}' new='{}'",
158                    existing.source_label(),
159                    source.source_label()
160                );
161            }
162            None => {
163                sources.insert(role, source);
164            }
165        }
166    }
167
168    /// Register one embedded wasm payload as the built-in install source for one role.
169    pub fn register_embedded_module_wasm(
170        role: CanisterRole,
171        source_label: impl Into<String>,
172        wasm_module: &'static [u8],
173    ) {
174        Self::register_embedded_module_source(
175            role,
176            ApprovedModuleSource::embedded(source_label.into(), wasm_module),
177        );
178    }
179
180    /// Register the control-plane resolver used by root-owned installation flows.
181    pub fn register_module_source_resolver(resolver: &'static dyn ModuleSourceResolver) {
182        let _ = MODULE_SOURCE_RESOLVER.set(resolver);
183    }
184
185    /// Return whether one embedded module source override has been registered.
186    #[must_use]
187    pub fn has_embedded_module_source(role: &CanisterRole) -> bool {
188        EMBEDDED_MODULE_SOURCES.get().is_some_and(|sources| {
189            let sources = sources
190                .lock()
191                .unwrap_or_else(std::sync::PoisonError::into_inner);
192            sources.contains_key(role)
193        })
194    }
195
196    /// Resolve the approved install source for one canister role through the registered driver.
197    pub(crate) async fn approved_module_source(
198        role: &CanisterRole,
199    ) -> Result<ApprovedModuleSource, InternalError> {
200        if let Some(source) = EMBEDDED_MODULE_SOURCES.get().and_then(|sources| {
201            let sources = sources
202                .lock()
203                .unwrap_or_else(std::sync::PoisonError::into_inner);
204            sources.get(role).cloned()
205        }) {
206            WasmStoreMetrics::record(
207                WasmStoreMetricOperation::SourceResolve,
208                WasmStoreMetricSource::Embedded,
209                WasmStoreMetricOutcome::Completed,
210                WasmStoreMetricReason::Ok,
211            );
212            return Ok(source);
213        }
214
215        let resolver = MODULE_SOURCE_RESOLVER.get().ok_or_else(|| {
216            WasmStoreMetrics::record(
217                WasmStoreMetricOperation::SourceResolve,
218                WasmStoreMetricSource::Resolver,
219                WasmStoreMetricOutcome::Failed,
220                WasmStoreMetricReason::InvalidState,
221            );
222            InternalError::workflow(
223                InternalErrorOrigin::Workflow,
224                "module source resolver is not registered; root/control-plane install flows are unavailable".to_string(),
225            )
226        })?;
227
228        match resolver.approved_module_source(role).await {
229            Ok(source) => {
230                WasmStoreMetrics::record(
231                    WasmStoreMetricOperation::SourceResolve,
232                    WasmStoreMetricSource::Resolver,
233                    WasmStoreMetricOutcome::Completed,
234                    WasmStoreMetricReason::Ok,
235                );
236                Ok(source)
237            }
238            Err(err) => {
239                WasmStoreMetrics::record(
240                    WasmStoreMetricOperation::SourceResolve,
241                    WasmStoreMetricSource::Resolver,
242                    WasmStoreMetricOutcome::Failed,
243                    WasmStoreMetricReason::StoreCall,
244                );
245                Err(InternalError::public(err))
246            }
247        }
248    }
249}