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#[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#[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 #[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 #[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 #[must_use]
84 pub fn source_label(&self) -> &str {
85 &self.source_label
86 }
87
88 #[must_use]
90 pub fn module_hash(&self) -> &[u8] {
91 &self.module_hash
92 }
93
94 #[must_use]
96 pub fn payload_size(&self) -> String {
97 byte_size(self.payload_size_bytes)
98 }
99
100 #[must_use]
102 pub const fn payload_size_bytes(&self) -> u64 {
103 self.payload_size_bytes
104 }
105
106 #[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 #[must_use]
117 pub const fn payload(&self) -> &ApprovedModulePayload {
118 &self.payload
119 }
120}
121
122#[async_trait]
127pub trait ModuleSourceResolver: Send + Sync {
128 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
139pub struct ModuleSourceRuntimeApi;
144
145impl ModuleSourceRuntimeApi {
146 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 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 pub fn register_module_source_resolver(resolver: &'static dyn ModuleSourceResolver) {
182 let _ = MODULE_SOURCE_RESOLVER.set(resolver);
183 }
184
185 #[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 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}