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 restore_group: u16,
55 pub verification_class: String,
56 pub verification_checks: Vec<VerificationCheck>,
57 pub snapshot_plan: SnapshotPlan,
58}
59
60impl DiscoveredMember {
61 fn into_fleet_member(self) -> FleetMember {
63 FleetMember {
64 role: self.role,
65 canister_id: self.canister_id,
66 parent_canister_id: self.parent_canister_id,
67 subnet_canister_id: self.subnet_canister_id,
68 controller_hint: self.controller_hint,
69 identity_mode: self.identity_mode,
70 restore_group: self.restore_group,
71 verification_class: self.verification_class,
72 verification_checks: self.verification_checks,
73 source_snapshot: SourceSnapshot {
74 snapshot_id: self.snapshot_plan.snapshot_id,
75 module_hash: self.snapshot_plan.module_hash,
76 wasm_hash: self.snapshot_plan.wasm_hash,
77 code_version: self.snapshot_plan.code_version,
78 artifact_path: self.snapshot_plan.artifact_path,
79 checksum_algorithm: self.snapshot_plan.checksum_algorithm,
80 checksum: self.snapshot_plan.checksum,
81 },
82 }
83 }
84}
85
86#[derive(Clone, Debug)]
91pub struct SnapshotPlan {
92 pub snapshot_id: String,
93 pub module_hash: Option<String>,
94 pub wasm_hash: Option<String>,
95 pub code_version: Option<String>,
96 pub artifact_path: String,
97 pub checksum_algorithm: String,
98 pub checksum: Option<String>,
99}
100
101#[derive(Clone, Debug, Eq, PartialEq)]
106pub struct RegistryEntry {
107 pub pid: String,
108 pub role: Option<String>,
109 pub kind: Option<String>,
110 pub parent_pid: Option<String>,
111}
112
113#[derive(Clone, Debug, Eq, PartialEq)]
118pub struct SnapshotTarget {
119 pub canister_id: String,
120 pub role: Option<String>,
121 pub parent_canister_id: Option<String>,
122}
123
124#[derive(Debug, ThisError)]
129pub enum DiscoveryError {
130 #[error("discovered fleet has no members")]
131 EmptyFleet,
132
133 #[error("duplicate discovered canister id {0}")]
134 DuplicateCanisterId(String),
135
136 #[error("discovered member {0} has no verification checks")]
137 MissingVerificationChecks(String),
138
139 #[error("registry JSON must be an array or {{\"Ok\": [...]}}")]
140 InvalidRegistryJsonShape,
141
142 #[error("registry JSON did not contain the requested canister {0}")]
143 CanisterNotInRegistry(String),
144
145 #[error(transparent)]
146 Json(#[from] serde_json::Error),
147}
148
149pub fn parse_registry_entries(registry_json: &str) -> Result<Vec<RegistryEntry>, DiscoveryError> {
151 let data = serde_json::from_str::<Value>(registry_json)?;
152 let entries = data
153 .get("Ok")
154 .and_then(Value::as_array)
155 .or_else(|| data.as_array())
156 .ok_or(DiscoveryError::InvalidRegistryJsonShape)?;
157
158 Ok(entries.iter().filter_map(parse_registry_entry).collect())
159}
160
161pub fn targets_from_registry(
163 registry: &[RegistryEntry],
164 canister_id: &str,
165 recursive: bool,
166) -> Result<Vec<SnapshotTarget>, DiscoveryError> {
167 let by_pid = registry
168 .iter()
169 .map(|entry| (entry.pid.as_str(), entry))
170 .collect::<BTreeMap<_, _>>();
171
172 let root = by_pid
173 .get(canister_id)
174 .ok_or_else(|| DiscoveryError::CanisterNotInRegistry(canister_id.to_string()))?;
175
176 let mut targets = Vec::new();
177 let mut seen = BTreeSet::new();
178 targets.push(SnapshotTarget {
179 canister_id: root.pid.clone(),
180 role: root.role.clone(),
181 parent_canister_id: root.parent_pid.clone(),
182 });
183 seen.insert(root.pid.clone());
184
185 let mut queue = VecDeque::from([root.pid.clone()]);
186 while let Some(parent) = queue.pop_front() {
187 for child in registry
188 .iter()
189 .filter(|entry| entry.parent_pid.as_deref() == Some(parent.as_str()))
190 {
191 if seen.insert(child.pid.clone()) {
192 targets.push(SnapshotTarget {
193 canister_id: child.pid.clone(),
194 role: child.role.clone(),
195 parent_canister_id: child.parent_pid.clone(),
196 });
197 if recursive {
198 queue.push_back(child.pid.clone());
199 }
200 }
201 }
202 }
203
204 Ok(targets)
205}
206
207fn parse_registry_entry(value: &Value) -> Option<RegistryEntry> {
209 let pid = value.get("pid").and_then(Value::as_str)?.to_string();
210 let role = value
211 .get("role")
212 .and_then(Value::as_str)
213 .map(str::to_string);
214 let parent_pid = value
215 .get("record")
216 .and_then(|record| record.get("parent_pid"))
217 .and_then(parse_optional_principal);
218 let kind = value
219 .get("kind")
220 .or_else(|| value.get("record").and_then(|record| record.get("kind")))
221 .and_then(Value::as_str)
222 .map(str::to_string);
223
224 Some(RegistryEntry {
225 pid,
226 role,
227 kind,
228 parent_pid,
229 })
230}
231
232fn parse_optional_principal(value: &Value) -> Option<String> {
234 if value.is_null() {
235 return None;
236 }
237 if let Some(text) = value.as_str() {
238 return Some(text.to_string());
239 }
240 value
241 .as_array()
242 .and_then(|items| items.first())
243 .and_then(Value::as_str)
244 .map(str::to_string)
245}
246
247fn validate_discovered_members(members: &[DiscoveredMember]) -> Result<(), DiscoveryError> {
249 if members.is_empty() {
250 return Err(DiscoveryError::EmptyFleet);
251 }
252
253 let mut canister_ids = BTreeSet::new();
254 for member in members {
255 if !canister_ids.insert(member.canister_id.clone()) {
256 return Err(DiscoveryError::DuplicateCanisterId(
257 member.canister_id.clone(),
258 ));
259 }
260 if member.verification_checks.is_empty() {
261 return Err(DiscoveryError::MissingVerificationChecks(
262 member.canister_id.clone(),
263 ));
264 }
265 }
266
267 Ok(())
268}
269
270#[cfg(test)]
271mod tests;