canic_core/ops/runtime/
install_source.rs

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

canic_core/ops/runtime/
install_source.rs
