1use crate::{
2 manifest::{FleetMember, FleetSection, IdentityMode, SourceSnapshot, VerificationCheck},
3 topology::{TopologyHasher, TopologyRecord},
4};
5use std::collections::BTreeSet;
6use thiserror::Error as ThisError;
7
8#[derive(Clone, Debug)]
13pub struct DiscoveredFleet {
14 pub topology_records: Vec<TopologyRecord>,
15 pub members: Vec<DiscoveredMember>,
16}
17
18impl DiscoveredFleet {
19 pub fn into_fleet_section(self) -> Result<FleetSection, DiscoveryError> {
21 validate_discovered_members(&self.members)?;
22
23 let topology_hash = TopologyHasher::hash(&self.topology_records);
24 let members = self
25 .members
26 .into_iter()
27 .map(DiscoveredMember::into_fleet_member)
28 .collect();
29
30 Ok(FleetSection {
31 topology_hash_algorithm: topology_hash.algorithm,
32 topology_hash_input: topology_hash.input,
33 discovery_topology_hash: topology_hash.hash.clone(),
34 pre_snapshot_topology_hash: topology_hash.hash.clone(),
35 topology_hash: topology_hash.hash,
36 members,
37 })
38 }
39}
40
41#[derive(Clone, Debug)]
46pub struct DiscoveredMember {
47 pub role: String,
48 pub canister_id: String,
49 pub parent_canister_id: Option<String>,
50 pub subnet_canister_id: Option<String>,
51 pub controller_hint: Option<String>,
52 pub identity_mode: IdentityMode,
53 pub restore_group: u16,
54 pub verification_class: String,
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 restore_group: self.restore_group,
70 verification_class: self.verification_class,
71 verification_checks: self.verification_checks,
72 source_snapshot: SourceSnapshot {
73 snapshot_id: self.snapshot_plan.snapshot_id,
74 module_hash: self.snapshot_plan.module_hash,
75 wasm_hash: self.snapshot_plan.wasm_hash,
76 code_version: self.snapshot_plan.code_version,
77 artifact_path: self.snapshot_plan.artifact_path,
78 checksum_algorithm: self.snapshot_plan.checksum_algorithm,
79 checksum: self.snapshot_plan.checksum,
80 },
81 }
82 }
83}
84
85#[derive(Clone, Debug)]
90pub struct SnapshotPlan {
91 pub snapshot_id: String,
92 pub module_hash: Option<String>,
93 pub wasm_hash: Option<String>,
94 pub code_version: Option<String>,
95 pub artifact_path: String,
96 pub checksum_algorithm: String,
97 pub checksum: Option<String>,
98}
99
100#[derive(Debug, ThisError)]
105pub enum DiscoveryError {
106 #[error("discovered fleet has no members")]
107 EmptyFleet,
108
109 #[error("duplicate discovered canister id {0}")]
110 DuplicateCanisterId(String),
111
112 #[error("discovered member {0} has no verification checks")]
113 MissingVerificationChecks(String),
114}
115
116fn validate_discovered_members(members: &[DiscoveredMember]) -> Result<(), DiscoveryError> {
118 if members.is_empty() {
119 return Err(DiscoveryError::EmptyFleet);
120 }
121
122 let mut canister_ids = BTreeSet::new();
123 for member in members {
124 if !canister_ids.insert(member.canister_id.clone()) {
125 return Err(DiscoveryError::DuplicateCanisterId(
126 member.canister_id.clone(),
127 ));
128 }
129 if member.verification_checks.is_empty() {
130 return Err(DiscoveryError::MissingVerificationChecks(
131 member.canister_id.clone(),
132 ));
133 }
134 }
135
136 Ok(())
137}
138
139#[cfg(test)]
140mod tests {
141 use super::*;
142 use candid::Principal;
143
144 const ROOT: Principal = Principal::from_slice(&[]);
145 const HASH: &str = "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef";
146
147 fn p(id: u8) -> Principal {
149 Principal::from_slice(&[id; 29])
150 }
151
152 #[test]
154 fn discovery_projects_to_valid_fleet_section() {
155 let fleet = DiscoveredFleet {
156 topology_records: vec![
157 topology_record(ROOT, None, "root"),
158 topology_record(p(2), Some(ROOT), "app"),
159 ],
160 members: vec![
161 discovered_member("root", ROOT.to_string(), None),
162 discovered_member("app", p(2).to_string(), Some(ROOT.to_string())),
163 ],
164 };
165
166 let section = fleet
167 .into_fleet_section()
168 .expect("discovery should project");
169
170 section.validate().expect("fleet section should validate");
171 assert_eq!(section.discovery_topology_hash, section.topology_hash);
172 assert_eq!(section.members.len(), 2);
173 }
174
175 #[test]
177 fn discovery_rejects_duplicate_canisters() {
178 let fleet = DiscoveredFleet {
179 topology_records: vec![topology_record(ROOT, None, "root")],
180 members: vec![
181 discovered_member("root", ROOT.to_string(), None),
182 discovered_member("root", ROOT.to_string(), None),
183 ],
184 };
185
186 let err = fleet
187 .into_fleet_section()
188 .expect_err("duplicate canisters should fail");
189
190 assert!(matches!(err, DiscoveryError::DuplicateCanisterId(_)));
191 }
192
193 #[test]
195 fn discovery_requires_verification_checks() {
196 let mut member = discovered_member("root", ROOT.to_string(), None);
197 member.verification_checks.clear();
198 let fleet = DiscoveredFleet {
199 topology_records: vec![topology_record(ROOT, None, "root")],
200 members: vec![member],
201 };
202
203 let err = fleet
204 .into_fleet_section()
205 .expect_err("missing verification should fail");
206
207 assert!(matches!(err, DiscoveryError::MissingVerificationChecks(_)));
208 }
209
210 fn topology_record(
212 pid: Principal,
213 parent_pid: Option<Principal>,
214 role: &str,
215 ) -> TopologyRecord {
216 TopologyRecord {
217 pid,
218 parent_pid,
219 role: role.to_string(),
220 module_hash: None,
221 }
222 }
223
224 fn discovered_member(
226 role: &str,
227 canister_id: String,
228 parent_canister_id: Option<String>,
229 ) -> DiscoveredMember {
230 DiscoveredMember {
231 role: role.to_string(),
232 canister_id,
233 parent_canister_id,
234 subnet_canister_id: None,
235 controller_hint: Some(ROOT.to_string()),
236 identity_mode: IdentityMode::Fixed,
237 restore_group: 1,
238 verification_class: "basic".to_string(),
239 verification_checks: vec![VerificationCheck {
240 kind: "call".to_string(),
241 method: Some("canic_ready".to_string()),
242 roles: Vec::new(),
243 }],
244 snapshot_plan: SnapshotPlan {
245 snapshot_id: format!("snap-{role}"),
246 module_hash: Some(HASH.to_string()),
247 wasm_hash: Some(HASH.to_string()),
248 code_version: Some("v0.30.0".to_string()),
249 artifact_path: format!("artifacts/{role}"),
250 checksum_algorithm: "sha256".to_string(),
251 checksum: Some(HASH.to_string()),
252 },
253 }
254 }
255}