1use crate::{
2 manifest::{FleetMember, FleetSection, IdentityMode, SourceSnapshot, VerificationCheck},
3 topology::{TopologyHasher, TopologyRecord},
4};
5use canic_cdk::utils::hash::hex_bytes;
6use serde_json::Value;
7use std::collections::{BTreeMap, BTreeSet, VecDeque};
8use thiserror::Error as ThisError;
9
10#[derive(Clone, Debug)]
15pub struct DiscoveredFleet {
16 pub topology_records: Vec<TopologyRecord>,
17 pub members: Vec<DiscoveredMember>,
18}
19
20impl DiscoveredFleet {
21 pub fn into_fleet_section(self) -> Result<FleetSection, DiscoveryError> {
23 validate_discovered_members(&self.members)?;
24
25 let topology_hash = TopologyHasher::hash(&self.topology_records);
26 let members = self
27 .members
28 .into_iter()
29 .map(DiscoveredMember::into_fleet_member)
30 .collect();
31
32 Ok(FleetSection {
33 topology_hash_algorithm: topology_hash.algorithm,
34 topology_hash_input: topology_hash.input,
35 discovery_topology_hash: topology_hash.hash.clone(),
36 pre_snapshot_topology_hash: topology_hash.hash.clone(),
37 topology_hash: topology_hash.hash,
38 members,
39 })
40 }
41}
42
43#[derive(Clone, Debug)]
48pub struct DiscoveredMember {
49 pub role: String,
50 pub canister_id: String,
51 pub parent_canister_id: Option<String>,
52 pub subnet_canister_id: Option<String>,
53 pub controller_hint: Option<String>,
54 pub identity_mode: IdentityMode,
55 pub verification_checks: Vec<VerificationCheck>,
56 pub snapshot_plan: SnapshotPlan,
57}
58
59impl DiscoveredMember {
60 fn into_fleet_member(self) -> FleetMember {
62 FleetMember {
63 role: self.role,
64 canister_id: self.canister_id,
65 parent_canister_id: self.parent_canister_id,
66 subnet_canister_id: self.subnet_canister_id,
67 controller_hint: self.controller_hint,
68 identity_mode: self.identity_mode,
69 verification_checks: self.verification_checks,
70 source_snapshot: SourceSnapshot {
71 snapshot_id: self.snapshot_plan.snapshot_id,
72 module_hash: self.snapshot_plan.module_hash,
73 code_version: self.snapshot_plan.code_version,
74 artifact_path: self.snapshot_plan.artifact_path,
75 checksum_algorithm: self.snapshot_plan.checksum_algorithm,
76 checksum: self.snapshot_plan.checksum,
77 },
78 }
79 }
80}
81
82#[derive(Clone, Debug)]
87pub struct SnapshotPlan {
88 pub snapshot_id: String,
89 pub module_hash: Option<String>,
90 pub code_version: Option<String>,
91 pub artifact_path: String,
92 pub checksum_algorithm: String,
93 pub checksum: Option<String>,
94}
95
96#[derive(Clone, Debug, Eq, PartialEq)]
101pub struct RegistryEntry {
102 pub pid: String,
103 pub role: Option<String>,
104 pub kind: Option<String>,
105 pub parent_pid: Option<String>,
106 pub module_hash: Option<String>,
107}
108
109#[derive(Clone, Debug, Eq, PartialEq)]
114pub struct SnapshotTarget {
115 pub canister_id: String,
116 pub role: Option<String>,
117 pub parent_canister_id: Option<String>,
118 pub module_hash: Option<String>,
119}
120
121#[derive(Debug, ThisError)]
126pub enum DiscoveryError {
127 #[error("discovered fleet has no members")]
128 EmptyFleet,
129
130 #[error("duplicate discovered canister id {0}")]
131 DuplicateCanisterId(String),
132
133 #[error("discovered member {0} has no verification checks")]
134 MissingVerificationChecks(String),
135
136 #[error("registry JSON must be an array or {{\"Ok\": [...]}}")]
137 InvalidRegistryJsonShape,
138
139 #[error("registry JSON did not contain the requested canister {0}")]
140 CanisterNotInRegistry(String),
141
142 #[error(transparent)]
143 Json(#[from] serde_json::Error),
144}
145
146pub fn parse_registry_entries(registry_json: &str) -> Result<Vec<RegistryEntry>, DiscoveryError> {
148 let data = serde_json::from_str::<Value>(registry_json)?;
149 let entries = data
150 .get("Ok")
151 .and_then(Value::as_array)
152 .or_else(|| data.as_array())
153 .ok_or(DiscoveryError::InvalidRegistryJsonShape)?;
154
155 Ok(entries.iter().filter_map(parse_registry_entry).collect())
156}
157
158pub fn targets_from_registry(
160 registry: &[RegistryEntry],
161 canister_id: &str,
162 recursive: bool,
163) -> Result<Vec<SnapshotTarget>, DiscoveryError> {
164 let by_pid = registry
165 .iter()
166 .map(|entry| (entry.pid.as_str(), entry))
167 .collect::<BTreeMap<_, _>>();
168
169 let root = by_pid
170 .get(canister_id)
171 .ok_or_else(|| DiscoveryError::CanisterNotInRegistry(canister_id.to_string()))?;
172
173 let mut targets = Vec::new();
174 let mut seen = BTreeSet::new();
175 targets.push(SnapshotTarget {
176 canister_id: root.pid.clone(),
177 role: root.role.clone(),
178 parent_canister_id: root.parent_pid.clone(),
179 module_hash: root.module_hash.clone(),
180 });
181 seen.insert(root.pid.clone());
182
183 let mut queue = VecDeque::from([root.pid.clone()]);
184 while let Some(parent) = queue.pop_front() {
185 for child in registry
186 .iter()
187 .filter(|entry| entry.parent_pid.as_deref() == Some(parent.as_str()))
188 {
189 if seen.insert(child.pid.clone()) {
190 targets.push(SnapshotTarget {
191 canister_id: child.pid.clone(),
192 role: child.role.clone(),
193 parent_canister_id: child.parent_pid.clone(),
194 module_hash: child.module_hash.clone(),
195 });
196 if recursive {
197 queue.push_back(child.pid.clone());
198 }
199 }
200 }
201 }
202
203 Ok(targets)
204}
205
206fn parse_registry_entry(value: &Value) -> Option<RegistryEntry> {
208 let pid = value.get("pid").and_then(Value::as_str)?.to_string();
209 let role = value
210 .get("role")
211 .and_then(Value::as_str)
212 .map(str::to_string);
213 let parent_pid = value
214 .get("record")
215 .and_then(|record| record.get("parent_pid"))
216 .and_then(parse_optional_principal);
217 let kind = value
218 .get("kind")
219 .or_else(|| value.get("record").and_then(|record| record.get("kind")))
220 .and_then(Value::as_str)
221 .map(str::to_string);
222 let module_hash = value
223 .get("record")
224 .and_then(|record| record.get("module_hash"))
225 .and_then(parse_module_hash);
226
227 Some(RegistryEntry {
228 pid,
229 role,
230 kind,
231 parent_pid,
232 module_hash,
233 })
234}
235
236fn parse_module_hash(value: &Value) -> Option<String> {
238 if value.is_null() {
239 return None;
240 }
241 if let Some(text) = value.as_str() {
242 return Some(text.to_string());
243 }
244 let bytes = value
245 .as_array()?
246 .iter()
247 .map(|item| {
248 let value = item.as_u64()?;
249 u8::try_from(value).ok()
250 })
251 .collect::<Option<Vec<_>>>()?;
252 Some(hex_bytes(bytes))
253}
254
255fn parse_optional_principal(value: &Value) -> Option<String> {
257 if value.is_null() {
258 return None;
259 }
260 if let Some(text) = value.as_str() {
261 return Some(text.to_string());
262 }
263 value
264 .as_array()
265 .and_then(|items| items.first())
266 .and_then(Value::as_str)
267 .map(str::to_string)
268}
269
270fn validate_discovered_members(members: &[DiscoveredMember]) -> Result<(), DiscoveryError> {
272 if members.is_empty() {
273 return Err(DiscoveryError::EmptyFleet);
274 }
275
276 let mut canister_ids = BTreeSet::new();
277 for member in members {
278 if !canister_ids.insert(member.canister_id.clone()) {
279 return Err(DiscoveryError::DuplicateCanisterId(
280 member.canister_id.clone(),
281 ));
282 }
283 if member.verification_checks.is_empty() {
284 return Err(DiscoveryError::MissingVerificationChecks(
285 member.canister_id.clone(),
286 ));
287 }
288 }
289
290 Ok(())
291}
292
293#[cfg(test)]
294mod tests;