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