1use anyhow::{Result, ensure};
17use aven_protocol::refusal::Bootstrap as Refusal;
18
19use super::exchange::{self, Link};
20use super::keys::ProtectedLocalKeyStore;
21use crate::db::Database;
22use crate::sync::shared_state::validated::{ProofCache, ValidatedSeed};
23use crate::sync::{
24 bootstrap_staging as staging,
25 seed_claim::{ClaimAuthentication, ClaimResult, Genesis, PublicationOutcome, Secret},
26};
27
28pub(crate) struct StageTimer {
30 stage: &'static str,
31 started: std::time::Instant,
32}
33
34impl StageTimer {
35 pub(crate) fn start(stage: &'static str) -> Self {
36 Self {
37 stage,
38 started: std::time::Instant::now(),
39 }
40 }
41}
42
43impl Drop for StageTimer {
44 fn drop(&mut self) {
45 tracing::debug!(
46 stage = self.stage,
47 elapsed_ms = self.started.elapsed().as_secs_f64() * 1000.0,
48 "setup stage finished"
49 );
50 }
51}
52
53pub use aven_protocol::wire::bootstrap::{
54 Envelope, Operation, PATH, REQUEST_LIMIT, RESPONSE_LIMIT, Reply,
55};
56
57impl From<staging::Status> for Reply {
58 fn from(value: staging::Status) -> Self {
59 match value {
60 staging::Status::Missing => Self::Missing,
61 staging::Status::Canceled => Self::Canceled,
62 staging::Status::Staging(status) => Self::Staging(status),
63 staging::Status::Published(outcome) => {
64 Self::Published(outcome.publication().record().to_vec())
65 }
66 }
67 }
68}
69
70pub struct Client {
72 link: Link,
73 endpoint: url::Url,
74}
75
76impl Client {
77 pub fn new(origin: &str, link: Link) -> Result<Self> {
78 Ok(Self {
79 link,
80 endpoint: super::origin::endpoint(origin, PATH)?,
81 })
82 }
83
84 pub async fn exchange(
85 &self,
86 genesis: &Genesis,
87 secret: &Secret,
88 operation: Operation,
89 ) -> Result<Reply> {
90 let claim = matches!(
91 &operation,
92 Operation::ClaimSetup { .. } | Operation::ClaimBearer { .. }
93 );
94 let bytes = serde_json::to_vec(&Envelope {
95 vault: genesis.context().vault_id,
96 genesis: genesis.commitment(),
97 operation,
98 })
99 .map_err(|_| anyhow::anyhow!("error bootstrap-request"))?;
100 ensure!(
101 bytes.len() <= REQUEST_LIMIT,
102 "error bootstrap-request-limit"
103 );
104 self.post(secret, exchange::json_content(), bytes, claim)
105 .await
106 }
107
108 pub async fn put_batch(
111 &self,
112 secret: &Secret,
113 header: &staging::batch::Header,
114 records: &[&[u8]],
115 ) -> Result<Reply> {
116 let bytes = staging::batch::encode(header, records)?;
117 let content_type = exchange::HttpHeader {
118 name: "content-type".into(),
119 value: staging::batch::CONTENT_TYPE.into(),
120 };
121 self.post(secret, content_type, bytes, false).await
122 }
123
124 async fn post(
125 &self,
126 secret: &Secret,
127 content_type: exchange::HttpHeader,
128 bytes: Vec<u8>,
129 claim: bool,
130 ) -> Result<Reply> {
131 let bytes = exchange::post(
132 &self.link,
133 &self.endpoint,
134 Some(secret),
135 content_type,
136 bytes,
137 RESPONSE_LIMIT,
138 )
139 .await
140 .map_err(|failure| match failure {
141 exchange::Failure::Hosting(hosting) => super::errors::hosting_error(hosting),
142 exchange::Failure::Network => {
143 anyhow::anyhow!("error bootstrap-network outcome-unknown")
144 }
145 exchange::Failure::SecureTransport => {
146 anyhow::anyhow!("error bootstrap-tls outcome-unknown")
147 }
148 exchange::Failure::Malformed => anyhow::anyhow!("error bootstrap-response"),
149 exchange::Failure::TooLarge => anyhow::anyhow!("error bootstrap-response-limit"),
150 exchange::Failure::RequestBodyLimit => {
151 anyhow::anyhow!("error sync-request-body-limit")
152 }
153 exchange::Failure::Refused { code, .. } => {
154 match Refusal::classify(code.as_deref(), claim) {
155 Refusal::Claimed => {
156 anyhow::anyhow!("error bootstrap-storage-already-claimed")
157 }
158 Refusal::SetupRejected => {
159 anyhow::anyhow!("error bootstrap-setup-invitation-rejected")
160 }
161 Refusal::SetupExpired => {
162 anyhow::anyhow!("error bootstrap-setup-invitation-expired")
163 }
164 Refusal::Quota => {
165 anyhow::anyhow!("error attachment-quota-exceeded")
166 }
167 Refusal::Unknown => anyhow::anyhow!("error bootstrap-refused outcome-unknown"),
168 }
169 }
170 })?;
171 serde_json::from_slice(&bytes).map_err(|_| anyhow::anyhow!("error bootstrap-response"))
172 }
173
174 pub async fn claim(
178 &self,
179 genesis: &Genesis,
180 authentication: ClaimAuthentication<'_>,
181 ) -> Result<()> {
182 let (secret, operation) = match authentication {
183 ClaimAuthentication::SetupSecret(secret) => (
184 secret,
185 Operation::ClaimSetup {
186 bytes: genesis.claim_bytes(),
187 },
188 ),
189 ClaimAuthentication::SeedBearer(secret) => (
190 secret,
191 Operation::ClaimBearer {
192 bytes: genesis.claim_bytes(),
193 },
194 ),
195 };
196 match self.exchange(genesis, secret, operation).await? {
197 Reply::Claimed {
198 vault,
199 claim,
200 genesis: commitment,
201 } => ClaimResult {
202 vault_id: vault,
203 claim_id: claim,
204 genesis_commitment: commitment,
205 }
206 .validate_pinned(genesis),
207 _ => anyhow::bail!("error bootstrap-response"),
208 }
209 }
210
211 pub async fn resume(
215 &self,
216 store: &ProtectedLocalKeyStore,
217 database: &Database,
218 ) -> Result<bool> {
219 self.resume_reporting(store, database, &|_, _| {}).await
220 }
221
222 pub async fn resume_reporting(
226 &self,
227 store: &ProtectedLocalKeyStore,
228 database: &Database,
229 uploaded: &(dyn Fn(u64, u64) + Sync),
230 ) -> Result<bool> {
231 self.resume_validated(store, database, None, uploaded).await
232 }
233
234 pub(crate) async fn resume_validated(
237 &self,
238 store: &ProtectedLocalKeyStore,
239 database: &Database,
240 proof: Option<ValidatedSeed>,
241 uploaded: &(dyn Fn(u64, u64) + Sync),
242 ) -> Result<bool> {
243 let mut proofs = ProofCache::new(proof);
244 let (seed, intent) = store.seed_resume_intent(database, &mut proofs).await?;
245 let publication = intent.publication(seed.genesis())?;
246 let binding = publication.binding();
247 let status = self
248 .exchange(
249 seed.genesis(),
250 seed.bearer(),
251 Operation::Status {
252 bootstrap: binding.bootstrap_id,
253 },
254 )
255 .await?;
256 let record = match status {
257 Reply::Published(record) => record,
258 Reply::Missing | Reply::Staging(_) => {
259 let upload = store
260 .seed_upload(database, &intent, &mut proofs)
261 .await?
262 .ok_or_else(|| anyhow::anyhow!("error bootstrap-outcome-missing"))?;
263 let _timer = StageTimer::start("upload");
264 let budget = upload.budget();
265 if let Reply::Staging(ref current) = status {
266 validate_staging(current, upload.slots(), binding, budget)?;
267 }
268 let staging = match status {
269 Reply::Missing => {
270 self.exchange(
271 seed.genesis(),
272 seed.bearer(),
273 Operation::Declare {
274 descriptor: upload.descriptor().to_vec(),
275 budget,
276 },
277 )
278 .await?
279 }
280 Reply::Staging(_) => {
281 self.exchange(
282 seed.genesis(),
283 seed.bearer(),
284 Operation::Ensure {
285 bootstrap: binding.bootstrap_id,
286 commitment: binding.descriptor_commitment,
287 },
288 )
289 .await?
290 }
291 _ => unreachable!(),
292 };
293 let Reply::Staging(staging) = staging else {
294 anyhow::bail!("error bootstrap-response");
295 };
296 validate_staging(&staging, upload.slots(), binding, budget)?;
297 let header = |records| staging::batch::Header {
298 vault: seed.genesis().context().vault_id,
299 genesis: seed.genesis().commitment(),
300 bootstrap: binding.bootstrap_id,
301 commitment: binding.descriptor_commitment,
302 records,
303 };
304 uploaded(0, budget.bytes);
305 let mut sent = 0_u64;
306 for batch in pack(upload.slots(), header) {
309 let records = upload.read(database, &batch).await?;
310 let refs: Vec<&[u8]> = records.iter().map(Vec::as_slice).collect();
311 let reply = self.put_batch(seed.bearer(), &header(batch), &refs).await?;
312 ensure!(
313 matches!(reply, Reply::Stored),
314 "error bootstrap-upload-refused"
315 );
316 sent += records
317 .iter()
318 .map(|record| record.len() as u64)
319 .sum::<u64>();
320 uploaded(sent, budget.bytes);
321 }
322 tracing::debug!(sent, total = budget.bytes, "uploaded frozen staged bytes");
323 match self
324 .exchange(
325 seed.genesis(),
326 seed.bearer(),
327 Operation::Publish {
328 bootstrap: binding.bootstrap_id,
329 commitment: binding.descriptor_commitment,
330 record: publication.record().to_vec(),
331 },
332 )
333 .await?
334 {
335 Reply::Published(record) => record,
336 _ => anyhow::bail!("error bootstrap-response"),
337 }
338 }
339 _ => anyhow::bail!("error bootstrap-candidate-unavailable"),
340 };
341 let outcome =
342 PublicationOutcome::from_response(seed.genesis(), intent.descriptor(), &record)?;
343 ensure!(
344 outcome.publication() == &publication,
345 "error bootstrap-outcome-mismatch"
346 );
347 let _timer = StageTimer::start("adopt");
348 store
349 .adopt_seed_publication_with(database, &outcome, &mut proofs)
350 .await
351 }
352}
353
354fn validate_staging(
357 status: &staging::StagingStatus,
358 slots: &[staging::batch::Slot],
359 binding: &crate::sync::seed_claim::PublicationBinding,
360 budget: staging::Budget,
361) -> Result<()> {
362 ensure!(
363 status.descriptor_commitment == binding.descriptor_commitment
364 && status.stream_id == binding.stream_id
365 && status.budget == budget,
366 "error bootstrap-status-mismatch"
367 );
368
369 let mut seen = Vec::new();
370 for listed in &status.components {
371 let expected = slots
372 .iter()
373 .filter(|slot| slot.component == listed.component)
374 .count();
375 ensure!(
376 expected > 0 && listed.chunks.len() == expected && !seen.contains(&listed.component),
377 "error bootstrap-status-mismatch"
378 );
379 seen.push(listed.component);
380 }
381
382 for component in [
385 staging::Component::DataCatalog,
386 staging::Component::PrefixCatalog,
387 staging::Component::ImageCatalog,
388 staging::Component::Manifest,
389 ] {
390 ensure!(seen.contains(&component), "error bootstrap-status-mismatch");
391 }
392
393 let catalog_complete = |component| {
394 status
395 .components
396 .iter()
397 .find(|listed| listed.component == component)
398 .is_some_and(|listed| {
399 listed
400 .chunks
401 .iter()
402 .all(|presence| *presence == staging::Presence::Verified)
403 })
404 };
405 let data_complete = catalog_complete(staging::Component::DataCatalog);
406 let images_complete = catalog_complete(staging::Component::ImageCatalog);
407 let listed = |component| seen.contains(&component);
408 ensure!(
409 listed(staging::Component::State) == data_complete,
410 "error bootstrap-status-mismatch"
411 );
412 let mut image_components = Vec::new();
413 for slot in slots {
414 if let staging::Component::Image(_) = slot.component
415 && !image_components.contains(&slot.component)
416 {
417 image_components.push(slot.component);
418 }
419 }
420 ensure!(
421 image_components
422 .iter()
423 .all(|component| listed(*component) == images_complete),
424 "error bootstrap-status-mismatch"
425 );
426 Ok(())
427}
428
429const UPLOAD_BATCH_PAYLOAD: u64 = staging::MAX_REQUEST_BYTES as u64;
432
433fn pack(
436 slots: &[staging::batch::Slot],
437 header: impl Fn(Vec<staging::batch::Slot>) -> staging::batch::Header,
438) -> Vec<Vec<staging::batch::Slot>> {
439 let mut batches: Vec<Vec<staging::batch::Slot>> = Vec::new();
440 let mut current: Vec<staging::batch::Slot> = Vec::new();
441 let mut payload = 0_u64;
442 for slot in slots {
443 let fits = !current.is_empty()
444 && current.len() < staging::batch::MAX_RECORDS
445 && payload + slot.len <= UPLOAD_BATCH_PAYLOAD
446 && current[0].component.is_catalog() == slot.component.is_catalog()
447 && staging::batch::header_len(&header(
448 current.iter().copied().chain([*slot]).collect(),
449 ))
450 .is_some();
451 if !fits && !current.is_empty() {
452 batches.push(std::mem::take(&mut current));
453 payload = 0;
454 }
455 current.push(*slot);
456 payload += slot.len;
457 }
458 if !current.is_empty() {
459 batches.push(current);
460 }
461 batches
462}
463
464#[cfg(any(test, feature = "test-support"))]
466pub fn components(
467 package: &crate::sync::bootstrap_format::Package,
468) -> Vec<(staging::Component, Vec<&[u8]>)> {
469 let mut out = Vec::new();
470 for (index, component) in [
471 staging::Component::DataCatalog,
472 staging::Component::PrefixCatalog,
473 staging::Component::ImageCatalog,
474 ]
475 .into_iter()
476 .enumerate()
477 {
478 out.push((
479 component,
480 package.catalogs[index].chunks(1_048_576).collect(),
481 ));
482 }
483 out.push((
484 staging::Component::Manifest,
485 package.manifest.iter().map(Vec::as_slice).collect(),
486 ));
487 out.push((
488 staging::Component::State,
489 package.state.iter().map(Vec::as_slice).collect(),
490 ));
491 for image in &package.images {
492 out.push((
493 staging::Component::Image(image.object_id),
494 image.records.iter().map(Vec::as_slice).collect(),
495 ));
496 }
497 out
498}
499
500#[cfg(test)]
501mod tests {
502 use super::*;
503
504 fn header(records: Vec<staging::batch::Slot>) -> staging::batch::Header {
505 staging::batch::Header {
506 vault: [1; 32],
507 genesis: [2; 32],
508 bootstrap: [3; 32],
509 commitment: [4; 32],
510 records,
511 }
512 }
513
514 #[test]
515 fn upload_batches_acknowledge_each_full_size_record() {
516 let slots: Vec<_> = (0..4)
517 .map(|index| staging::batch::Slot {
518 component: staging::Component::State,
519 index,
520 len: staging::MAX_REQUEST_BYTES as u64,
521 })
522 .collect();
523 let batches = pack(&slots, header);
524 assert!(batches.iter().all(|batch| batch.len() == 1));
525 assert_eq!(batches.concat(), slots);
526 }
527
528 #[test]
529 fn upload_batches_group_small_records_without_mixing_catalogs() {
530 let slots: Vec<_> = (0..10)
531 .map(|index| staging::batch::Slot {
532 component: if index < 2 {
533 staging::Component::DataCatalog
534 } else {
535 staging::Component::State
536 },
537 index,
538 len: UPLOAD_BATCH_PAYLOAD / 4,
539 })
540 .collect();
541 let batches = pack(&slots, header);
542 assert_eq!(batches.iter().map(Vec::len).collect::<Vec<_>>(), [2, 4, 4]);
543 assert_eq!(batches.concat(), slots);
544 for batch in batches {
545 let records: Vec<_> = batch
546 .iter()
547 .map(|slot| vec![0; slot.len as usize])
548 .collect();
549 let refs: Vec<_> = records.iter().map(Vec::as_slice).collect();
550 assert!(staging::batch::encode(&header(batch), &refs).is_ok());
551 }
552 }
553}