1use std::fs::File;
12use std::os::unix::fs::FileExt;
13use std::sync::Arc;
14
15use anyhow::{Result, anyhow};
16use arrow::array::{
17 BooleanArray, BooleanBuilder, FixedSizeBinaryArray, FixedSizeBinaryBuilder, StringArray,
18 StringBuilder, UInt32Array, UInt32Builder, UInt64Array, UInt64Builder,
19};
20use arrow::datatypes::Schema;
21use arrow::ipc::writer::StreamWriter;
22use arrow::record_batch::RecordBatch;
23
24use crate::index::{
25 ChunkLoc, LOOKUP_MODULE, META_MODULE, MULTI_INDEX_MAGIC, ManifestEntry, RESERVED_PKG_TYPE,
26 TRIE_MODULE, is_reserved_module, lookup_schema, write_manifest_bytes,
27};
28use crate::meta_index::{MetaTable, build_meta_batch, meta_schema};
29#[cfg(feature = "sign")]
30use crate::index::{SIGN_ARCHIVE_MODULE, SIGN_ARTIFACTS_MODULE};
31
32#[derive(Debug, Clone)]
34pub struct GroupKey {
35 pub pkg_type: i8,
36 pub repo: String,
37 pub module_name: String,
38}
39
40#[derive(Clone)]
54pub enum ReservedPayload {
55 Raw(Vec<u8>),
56 Arrow { schema: Arc<Schema>, batches: Vec<RecordBatch> },
57}
58
59#[derive(Clone)]
64pub struct ReservedSection {
65 pub module_name: String,
66 pub payload: ReservedPayload,
67}
68
69impl std::fmt::Debug for ReservedSection {
72 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
73 match &self.payload {
74 ReservedPayload::Raw(b) => f
75 .debug_struct("ReservedSection")
76 .field("module_name", &self.module_name)
77 .field("raw_bytes", &b.len())
78 .finish(),
79 ReservedPayload::Arrow { batches, .. } => f
80 .debug_struct("ReservedSection")
81 .field("module_name", &self.module_name)
82 .field("rows", &batches.iter().map(|b| b.num_rows()).sum::<usize>())
83 .finish(),
84 }
85 }
86}
87
88impl ReservedSection {
89 pub fn raw(module_name: impl Into<String>, bytes: Vec<u8>) -> Self {
90 Self { module_name: module_name.into(), payload: ReservedPayload::Raw(bytes) }
91 }
92
93 pub fn arrow(
94 module_name: impl Into<String>,
95 schema: Arc<Schema>,
96 batches: Vec<RecordBatch>,
97 ) -> Self {
98 Self {
99 module_name: module_name.into(),
100 payload: ReservedPayload::Arrow { schema, batches },
101 }
102 }
103}
104
105pub struct LookupView<'a> {
113 paths: &'a [String],
114 locs: &'a [ChunkLoc],
115 order: &'a [usize],
116}
117
118impl<'a> LookupView<'a> {
119 pub fn len(&self) -> usize {
121 self.order.len()
122 }
123
124 pub fn is_empty(&self) -> bool {
125 self.order.is_empty()
126 }
127
128 pub fn path(&self, row: usize) -> &'a str {
130 &self.paths[self.order[row]]
131 }
132
133 pub fn loc(&self, row: usize) -> &'a ChunkLoc {
135 &self.locs[self.order[row]]
136 }
137
138 pub fn first_rows(&self) -> Vec<(&'a str, u64)> {
141 let mut out: Vec<(&'a str, u64)> = Vec::new();
142 let mut prev: Option<&str> = None;
143 for row in 0..self.order.len() {
144 let p = self.path(row);
145 if prev != Some(p) {
146 out.push((p, row as u64));
147 prev = Some(p);
148 }
149 }
150 out
151 }
152}
153
154pub type ReservedSectionBuilder =
160 Box<dyn FnOnce(&LookupView<'_>) -> Result<Vec<ReservedSection>> + Send>;
161
162pub trait ArchiveMetaSink {
167 fn push_subindex(
171 &mut self,
172 schema: &Schema,
173 batches: &[RecordBatch],
174 key: GroupKey,
175 ) -> Result<()>;
176
177 fn finish(self: Box<Self>) -> Result<u64>;
179}
180
181pub type MetaSinkFactory = Box<dyn FnOnce(Arc<File>, u64) -> Box<dyn ArchiveMetaSink> + Send>;
192
193pub struct ArrowIpcSink {
197 file: Arc<File>,
198 cursor: u64,
199 entries: Vec<ManifestEntry>,
200 lookup_paths: Vec<String>,
203 lookup_locs: Vec<ChunkLoc>,
204 meta: Option<MetaTable>,
210 reserved_builder: Option<ReservedSectionBuilder>,
214 #[cfg(feature = "sign")]
219 signer: Option<Box<dyn crate::sign::ArchiveSigner + Send>>,
220}
221
222impl ArrowIpcSink {
223 pub fn new(file: Arc<File>, blob_end_offset: u64) -> Self {
226 Self {
227 file,
228 cursor: blob_end_offset,
229 entries: Vec::new(),
230 lookup_paths: Vec::new(),
231 lookup_locs: Vec::new(),
232 meta: None,
233 reserved_builder: None,
234 #[cfg(feature = "sign")]
235 signer: None,
236 }
237 }
238
239 pub fn with_reserved_builder(mut self, builder: ReservedSectionBuilder) -> Self {
248 self.reserved_builder = Some(builder);
249 self
250 }
251
252 pub fn with_meta(mut self, meta: MetaTable) -> Self {
259 self.meta = Some(meta);
260 self
261 }
262
263 #[cfg(feature = "sign")]
266 pub fn with_signer(mut self, signer: Box<dyn crate::sign::ArchiveSigner + Send>) -> Self {
267 self.signer = Some(signer);
268 self
269 }
270
271 fn accumulate_lookup(&mut self, batch: &RecordBatch) {
275 let cols = (|| {
276 Some((
277 batch.column_by_name("relative_path")?.as_any().downcast_ref::<StringArray>()?,
278 batch.column_by_name("chunk_seq")?.as_any().downcast_ref::<UInt32Array>()?,
279 batch.column_by_name("fdata_offset")?.as_any().downcast_ref::<UInt64Array>()?,
280 batch.column_by_name("compressed")?.as_any().downcast_ref::<BooleanArray>()?,
281 batch.column_by_name("uncompressed_size")?.as_any().downcast_ref::<UInt64Array>()?,
282 batch.column_by_name("blob_offset")?.as_any().downcast_ref::<UInt64Array>()?,
283 batch.column_by_name("blob_size")?.as_any().downcast_ref::<UInt64Array>()?,
284 batch.column_by_name("checksum")?.as_any().downcast_ref::<FixedSizeBinaryArray>()?,
285 ))
286 })();
287 let Some((paths, chunk_seq, fdata, compressed, usz, blob_off, blob_sz, checksum)) = cols
288 else { return; };
289 for i in 0..batch.num_rows() {
290 let mut ck = [0u8; 32];
291 ck.copy_from_slice(checksum.value(i));
292 self.lookup_paths.push(paths.value(i).to_string());
293 self.lookup_locs.push(ChunkLoc {
294 chunk_seq: chunk_seq.value(i),
295 fdata_offset: fdata.value(i),
296 blob_offset: blob_off.value(i),
297 blob_size: blob_sz.value(i),
298 uncompressed_size: usz.value(i),
299 compressed: compressed.value(i),
300 checksum: ck,
301 });
302 }
303 }
304
305 fn lookup_order(&self) -> Vec<usize> {
312 let n = self.lookup_paths.len();
313 let mut order: Vec<usize> = (0..n).collect();
314 order.sort_by(|&a, &b| {
315 self.lookup_paths[a].cmp(&self.lookup_paths[b])
316 .then(self.lookup_locs[a].chunk_seq.cmp(&self.lookup_locs[b].chunk_seq))
317 });
318 order
319 }
320
321 fn write_reserved_sections(&mut self, order: &[usize]) -> Result<()> {
326 let Some(builder) = self.reserved_builder.take() else {
327 return Ok(());
328 };
329 let sections = {
330 let view = LookupView {
331 paths: &self.lookup_paths,
332 locs: &self.lookup_locs,
333 order,
334 };
335 builder(&view)?
336 };
337 for section in sections {
338 anyhow::ensure!(
339 is_reserved_module(§ion.module_name),
340 "module '{}' is not a reserved module name; a non-reserved extra \
341 section would be merged into the data index and corrupt list/decompress",
342 section.module_name,
343 );
344 let key = GroupKey {
345 pkg_type: RESERVED_PKG_TYPE,
346 repo: String::new(),
347 module_name: section.module_name,
348 };
349 match section.payload {
350 ReservedPayload::Raw(bytes) => self.write_raw_section(&bytes, key)?,
351 ReservedPayload::Arrow { schema, batches } => {
352 self.push_subindex(schema.as_ref(), &batches, key)?
353 }
354 }
355 }
356 Ok(())
357 }
358
359 fn write_lookup_and_trie(&mut self, order: &[usize]) -> Result<()> {
362 let n = self.lookup_paths.len();
363
364 let mut path_b = StringBuilder::with_capacity(n, n * 16);
366 let mut seq_b = UInt32Builder::with_capacity(n);
367 let mut fdata_b = UInt64Builder::with_capacity(n);
368 let mut comp_b = BooleanBuilder::with_capacity(n);
369 let mut usz_b = UInt64Builder::with_capacity(n);
370 let mut boff_b = UInt64Builder::with_capacity(n);
371 let mut bsz_b = UInt64Builder::with_capacity(n);
372 let mut ck_b = FixedSizeBinaryBuilder::with_capacity(n, 32);
373 for &i in order {
374 let loc = &self.lookup_locs[i];
375 path_b.append_value(&self.lookup_paths[i]);
376 seq_b.append_value(loc.chunk_seq);
377 fdata_b.append_value(loc.fdata_offset);
378 comp_b.append_value(loc.compressed);
379 usz_b.append_value(loc.uncompressed_size);
380 boff_b.append_value(loc.blob_offset);
381 bsz_b.append_value(loc.blob_size);
382 ck_b.append_value(loc.checksum).expect("checksum is 32 bytes");
383 }
384 let schema = lookup_schema();
385 let batch = RecordBatch::try_new(
386 schema.clone(),
387 vec![
388 Arc::new(path_b.finish()),
389 Arc::new(seq_b.finish()),
390 Arc::new(fdata_b.finish()),
391 Arc::new(comp_b.finish()),
392 Arc::new(usz_b.finish()),
393 Arc::new(boff_b.finish()),
394 Arc::new(bsz_b.finish()),
395 Arc::new(ck_b.finish()),
396 ],
397 )?;
398 self.push_subindex(&schema, &[batch], GroupKey {
399 pkg_type: RESERVED_PKG_TYPE,
400 repo: String::new(),
401 module_name: LOOKUP_MODULE.to_string(),
402 })?;
403
404 let mut builder = fst::MapBuilder::memory();
408 let mut prev: Option<&str> = None;
409 for (sorted_idx, &orig) in order.iter().enumerate() {
410 let p = self.lookup_paths[orig].as_str();
411 if prev != Some(p) {
412 builder.insert(p.as_bytes(), sorted_idx as u64)
413 .map_err(|e| anyhow!("trie insert: {e}"))?;
414 prev = Some(p);
415 }
416 }
417 let trie_bytes = builder.into_inner().map_err(|e| anyhow!("trie finish: {e}"))?;
418 self.write_raw_section(&trie_bytes, GroupKey {
419 pkg_type: RESERVED_PKG_TYPE,
420 repo: String::new(),
421 module_name: TRIE_MODULE.to_string(),
422 })
423 }
424
425 #[cfg(feature = "sign")]
430 fn write_signatures(&mut self) -> Result<()> {
431 use std::collections::BTreeMap;
432 let Some(signer) = self.signer.take() else {
433 return Ok(());
434 };
435
436 let (file_digests, artifact_paths, artifact_cms): (
439 Vec<(String, [u8; 32])>,
440 Vec<String>,
441 Vec<Vec<u8>>,
442 ) = {
443 let mut by_path: BTreeMap<&str, Vec<(u32, &[u8; 32])>> = BTreeMap::new();
444 for (p, loc) in self.lookup_paths.iter().zip(self.lookup_locs.iter()) {
445 by_path.entry(p.as_str()).or_default().push((loc.chunk_seq, &loc.checksum));
446 }
447 let mut digs = Vec::with_capacity(by_path.len());
448 let mut paths = Vec::with_capacity(by_path.len());
449 let mut cmss = Vec::with_capacity(by_path.len());
450 for (path, mut chunks) in by_path {
451 chunks.sort_by_key(|(seq, _)| *seq);
452 let n = chunks.len();
453 let digest = crate::sign::file_digest_from_parts(
454 path,
455 chunks.iter().map(|(s, c)| (*s, *c)),
456 n,
457 );
458 let cms = signer.sign_digest(&digest)?;
459 digs.push((path.to_string(), digest));
460 paths.push(path.to_string());
461 cmss.push(cms);
462 }
463 (digs, paths, cmss)
464 };
465
466 let footer = crate::index::IndexFooter::Multi { manifest_offset: 0 };
468 let root = crate::sign::archive_root(&file_digests, &footer);
469 let archive_cms = signer.sign_digest(&root)?;
470
471 let artifacts_bytes = serialize_artifact_signatures(&artifact_paths, &artifact_cms)?;
472 self.write_raw_section(
473 &artifacts_bytes,
474 GroupKey {
475 pkg_type: RESERVED_PKG_TYPE,
476 repo: String::new(),
477 module_name: SIGN_ARTIFACTS_MODULE.to_string(),
478 },
479 )?;
480 self.write_raw_section(
481 &archive_cms,
482 GroupKey {
483 pkg_type: RESERVED_PKG_TYPE,
484 repo: String::new(),
485 module_name: SIGN_ARCHIVE_MODULE.to_string(),
486 },
487 )?;
488 Ok(())
489 }
490
491 fn write_meta_subindex(&mut self) -> Result<()> {
498 let Some(table) = self.meta.take() else {
499 return Ok(());
500 };
501 let batch = build_meta_batch(&table)?;
502 let schema = meta_schema();
503 self.push_subindex(schema.as_ref(), &[batch], GroupKey {
504 pkg_type: RESERVED_PKG_TYPE,
505 repo: String::new(),
506 module_name: META_MODULE.to_string(),
507 })
508 }
509
510 fn write_raw_section(&mut self, bytes: &[u8], key: GroupKey) -> Result<()> {
513 let start = self.cursor;
514 self.file.write_all_at(bytes, start)?;
515 self.cursor += bytes.len() as u64;
516 self.entries.push(ManifestEntry {
517 pkg_type: key.pkg_type,
518 repo: key.repo,
519 module_name: key.module_name,
520 index_offset: start,
521 index_len: bytes.len() as u64,
522 row_count: 0,
523 });
524 Ok(())
525 }
526}
527
528impl ArchiveMetaSink for ArrowIpcSink {
529 fn push_subindex(
530 &mut self,
531 schema: &Schema,
532 batches: &[RecordBatch],
533 key: GroupKey,
534 ) -> Result<()> {
535 let sub_start = self.cursor;
536 let mut sub_bytes: Vec<u8> = Vec::new();
537 let mut sw = StreamWriter::try_new(&mut sub_bytes, schema)
538 .map_err(|e| anyhow!("sub-index writer: {e}"))?;
539 let mut row_count = 0u64;
540 for batch in batches {
541 row_count += batch.num_rows() as u64;
542 sw.write(batch).map_err(|e| anyhow!("sub-index write: {e}"))?;
543 }
544 sw.finish().map_err(|e| anyhow!("sub-index finish: {e}"))?;
545
546 if !is_reserved_module(&key.module_name) {
555 for batch in batches {
556 self.accumulate_lookup(batch);
557 }
558 }
559
560 let sub_len = sub_bytes.len() as u64;
561 self.file.write_all_at(&sub_bytes, sub_start)?;
562 self.cursor += sub_len;
563
564 self.entries.push(ManifestEntry {
565 pkg_type: key.pkg_type,
566 repo: key.repo,
567 module_name: key.module_name,
568 index_offset: sub_start,
569 index_len: sub_len,
570 row_count,
571 });
572 Ok(())
573 }
574
575 fn finish(mut self: Box<Self>) -> Result<u64> {
576 let order = self.lookup_order();
579 self.write_lookup_and_trie(&order)?;
580
581 self.write_reserved_sections(&order)?;
585
586 self.write_meta_subindex()?;
590
591 #[cfg(feature = "sign")]
595 self.write_signatures()?;
596
597 let manifest_offset = self.cursor;
598 let manifest_bytes =
599 write_manifest_bytes(&self.entries).map_err(|e| anyhow!("manifest: {e}"))?;
600 self.file.write_all_at(&manifest_bytes, manifest_offset)?;
601
602 let after = manifest_offset + manifest_bytes.len() as u64;
603 self.file.write_all_at(&MULTI_INDEX_MAGIC, after)?;
604 self.file.write_all_at(
605 &manifest_offset.to_le_bytes(),
606 after + MULTI_INDEX_MAGIC.len() as u64,
607 )?;
608 self.file.sync_all()?;
609
610 Ok(after + MULTI_INDEX_MAGIC.len() as u64 + 8)
611 }
612}
613
614#[cfg(feature = "sign")]
618fn serialize_artifact_signatures(paths: &[String], cms: &[Vec<u8>]) -> Result<Vec<u8>> {
619 use arrow::array::BinaryBuilder;
620 use arrow::datatypes::{DataType, Field, Schema};
621
622 let n = paths.len();
623 let schema = Arc::new(Schema::new(vec![
624 Field::new("relative_path", DataType::Utf8, false),
625 Field::new("cms", DataType::Binary, false),
626 ]));
627 let mut path_b = StringBuilder::with_capacity(n, n * 32);
628 let mut cms_b = BinaryBuilder::with_capacity(n, n * 512);
629 for (p, c) in paths.iter().zip(cms.iter()) {
630 path_b.append_value(p);
631 cms_b.append_value(c);
632 }
633 let batch = RecordBatch::try_new(
634 schema.clone(),
635 vec![Arc::new(path_b.finish()), Arc::new(cms_b.finish())],
636 )?;
637 let mut buf = Vec::new();
638 {
639 let mut w = StreamWriter::try_new(&mut buf, &schema)
640 .map_err(|e| anyhow!("artifact-sig writer: {e}"))?;
641 w.write(&batch).map_err(|e| anyhow!("artifact-sig write: {e}"))?;
642 w.finish().map_err(|e| anyhow!("artifact-sig finish: {e}"))?;
643 }
644 Ok(buf)
645}