Skip to main content

canic_backup/discovery/
mod.rs

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///
10/// DiscoveredFleet
11///
12
13#[derive(Clone, Debug)]
14pub struct DiscoveredFleet {
15    pub topology_records: Vec<TopologyRecord>,
16    pub members: Vec<DiscoveredMember>,
17}
18
19impl DiscoveredFleet {
20    /// Convert discovered topology and member policy into a manifest fleet section.
21    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///
43/// DiscoveredMember
44///
45
46#[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    /// Project this discovery member into the manifest restore contract.
62    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///
87/// SnapshotPlan
88///
89
90#[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///
102/// RegistryEntry
103///
104
105#[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///
114/// SnapshotTarget
115///
116
117#[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///
125/// DiscoveryError
126///
127
128#[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
149/// Parse the `dfx --output json` subnet registry shape.
150pub 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
161/// Resolve selected target and children from registry entries.
162pub 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
207// Parse one registry entry from dfx JSON.
208fn 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
232// Parse optional principal JSON emitted as null, string, or optional vector form.
233fn 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
247// Validate discovery output before building a manifest projection.
248fn 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;