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
40pub enum ReservedPayload {
47 Raw(Vec<u8>),
48 Arrow { schema: Arc<Schema>, batches: Vec<RecordBatch> },
49}
50
51pub struct ReservedSection {
56 pub module_name: String,
57 pub payload: ReservedPayload,
58}
59
60impl std::fmt::Debug for ReservedSection {
63 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
64 match &self.payload {
65 ReservedPayload::Raw(b) => f
66 .debug_struct("ReservedSection")
67 .field("module_name", &self.module_name)
68 .field("raw_bytes", &b.len())
69 .finish(),
70 ReservedPayload::Arrow { batches, .. } => f
71 .debug_struct("ReservedSection")
72 .field("module_name", &self.module_name)
73 .field("rows", &batches.iter().map(|b| b.num_rows()).sum::<usize>())
74 .finish(),
75 }
76 }
77}
78
79impl ReservedSection {
80 pub fn raw(module_name: impl Into<String>, bytes: Vec<u8>) -> Self {
81 Self { module_name: module_name.into(), payload: ReservedPayload::Raw(bytes) }
82 }
83
84 pub fn arrow(
85 module_name: impl Into<String>,
86 schema: Arc<Schema>,
87 batches: Vec<RecordBatch>,
88 ) -> Self {
89 Self {
90 module_name: module_name.into(),
91 payload: ReservedPayload::Arrow { schema, batches },
92 }
93 }
94}
95
96pub struct LookupView<'a> {
104 paths: &'a [String],
105 locs: &'a [ChunkLoc],
106 order: &'a [usize],
107}
108
109impl<'a> LookupView<'a> {
110 pub fn len(&self) -> usize {
112 self.order.len()
113 }
114
115 pub fn is_empty(&self) -> bool {
116 self.order.is_empty()
117 }
118
119 pub fn path(&self, row: usize) -> &'a str {
121 &self.paths[self.order[row]]
122 }
123
124 pub fn loc(&self, row: usize) -> &'a ChunkLoc {
126 &self.locs[self.order[row]]
127 }
128
129 pub fn first_rows(&self) -> Vec<(&'a str, u64)> {
132 let mut out: Vec<(&'a str, u64)> = Vec::new();
133 let mut prev: Option<&str> = None;
134 for row in 0..self.order.len() {
135 let p = self.path(row);
136 if prev != Some(p) {
137 out.push((p, row as u64));
138 prev = Some(p);
139 }
140 }
141 out
142 }
143}
144
145pub type ReservedSectionBuilder =
151 Box<dyn FnOnce(&LookupView<'_>) -> Result<Vec<ReservedSection>> + Send>;
152
153pub trait ArchiveMetaSink {
158 fn push_subindex(
162 &mut self,
163 schema: &Schema,
164 batches: &[RecordBatch],
165 key: GroupKey,
166 ) -> Result<()>;
167
168 fn finish(self: Box<Self>) -> Result<u64>;
170}
171
172pub type MetaSinkFactory = Box<dyn FnOnce(Arc<File>, u64) -> Box<dyn ArchiveMetaSink> + Send>;
183
184pub struct ArrowIpcSink {
188 file: Arc<File>,
189 cursor: u64,
190 entries: Vec<ManifestEntry>,
191 lookup_paths: Vec<String>,
194 lookup_locs: Vec<ChunkLoc>,
195 meta: Option<MetaTable>,
201 reserved_builder: Option<ReservedSectionBuilder>,
205 #[cfg(feature = "sign")]
210 signer: Option<Box<dyn crate::sign::ArchiveSigner + Send>>,
211}
212
213impl ArrowIpcSink {
214 pub fn new(file: Arc<File>, blob_end_offset: u64) -> Self {
217 Self {
218 file,
219 cursor: blob_end_offset,
220 entries: Vec::new(),
221 lookup_paths: Vec::new(),
222 lookup_locs: Vec::new(),
223 meta: None,
224 reserved_builder: None,
225 #[cfg(feature = "sign")]
226 signer: None,
227 }
228 }
229
230 pub fn with_reserved_builder(mut self, builder: ReservedSectionBuilder) -> Self {
239 self.reserved_builder = Some(builder);
240 self
241 }
242
243 pub fn with_meta(mut self, meta: MetaTable) -> Self {
250 self.meta = Some(meta);
251 self
252 }
253
254 #[cfg(feature = "sign")]
257 pub fn with_signer(mut self, signer: Box<dyn crate::sign::ArchiveSigner + Send>) -> Self {
258 self.signer = Some(signer);
259 self
260 }
261
262 fn accumulate_lookup(&mut self, batch: &RecordBatch) {
266 let cols = (|| {
267 Some((
268 batch.column_by_name("relative_path")?.as_any().downcast_ref::<StringArray>()?,
269 batch.column_by_name("chunk_seq")?.as_any().downcast_ref::<UInt32Array>()?,
270 batch.column_by_name("fdata_offset")?.as_any().downcast_ref::<UInt64Array>()?,
271 batch.column_by_name("compressed")?.as_any().downcast_ref::<BooleanArray>()?,
272 batch.column_by_name("uncompressed_size")?.as_any().downcast_ref::<UInt64Array>()?,
273 batch.column_by_name("blob_offset")?.as_any().downcast_ref::<UInt64Array>()?,
274 batch.column_by_name("blob_size")?.as_any().downcast_ref::<UInt64Array>()?,
275 batch.column_by_name("checksum")?.as_any().downcast_ref::<FixedSizeBinaryArray>()?,
276 ))
277 })();
278 let Some((paths, chunk_seq, fdata, compressed, usz, blob_off, blob_sz, checksum)) = cols
279 else { return; };
280 for i in 0..batch.num_rows() {
281 let mut ck = [0u8; 32];
282 ck.copy_from_slice(checksum.value(i));
283 self.lookup_paths.push(paths.value(i).to_string());
284 self.lookup_locs.push(ChunkLoc {
285 chunk_seq: chunk_seq.value(i),
286 fdata_offset: fdata.value(i),
287 blob_offset: blob_off.value(i),
288 blob_size: blob_sz.value(i),
289 uncompressed_size: usz.value(i),
290 compressed: compressed.value(i),
291 checksum: ck,
292 });
293 }
294 }
295
296 fn lookup_order(&self) -> Vec<usize> {
303 let n = self.lookup_paths.len();
304 let mut order: Vec<usize> = (0..n).collect();
305 order.sort_by(|&a, &b| {
306 self.lookup_paths[a].cmp(&self.lookup_paths[b])
307 .then(self.lookup_locs[a].chunk_seq.cmp(&self.lookup_locs[b].chunk_seq))
308 });
309 order
310 }
311
312 fn write_reserved_sections(&mut self, order: &[usize]) -> Result<()> {
317 let Some(builder) = self.reserved_builder.take() else {
318 return Ok(());
319 };
320 let sections = {
321 let view = LookupView {
322 paths: &self.lookup_paths,
323 locs: &self.lookup_locs,
324 order,
325 };
326 builder(&view)?
327 };
328 for section in sections {
329 anyhow::ensure!(
330 is_reserved_module(§ion.module_name),
331 "module '{}' is not a reserved module name; a non-reserved extra \
332 section would be merged into the data index and corrupt list/decompress",
333 section.module_name,
334 );
335 let key = GroupKey {
336 pkg_type: RESERVED_PKG_TYPE,
337 repo: String::new(),
338 module_name: section.module_name,
339 };
340 match section.payload {
341 ReservedPayload::Raw(bytes) => self.write_raw_section(&bytes, key)?,
342 ReservedPayload::Arrow { schema, batches } => {
343 self.push_subindex(schema.as_ref(), &batches, key)?
344 }
345 }
346 }
347 Ok(())
348 }
349
350 fn write_lookup_and_trie(&mut self, order: &[usize]) -> Result<()> {
353 let n = self.lookup_paths.len();
354
355 let mut path_b = StringBuilder::with_capacity(n, n * 16);
357 let mut seq_b = UInt32Builder::with_capacity(n);
358 let mut fdata_b = UInt64Builder::with_capacity(n);
359 let mut comp_b = BooleanBuilder::with_capacity(n);
360 let mut usz_b = UInt64Builder::with_capacity(n);
361 let mut boff_b = UInt64Builder::with_capacity(n);
362 let mut bsz_b = UInt64Builder::with_capacity(n);
363 let mut ck_b = FixedSizeBinaryBuilder::with_capacity(n, 32);
364 for &i in order {
365 let loc = &self.lookup_locs[i];
366 path_b.append_value(&self.lookup_paths[i]);
367 seq_b.append_value(loc.chunk_seq);
368 fdata_b.append_value(loc.fdata_offset);
369 comp_b.append_value(loc.compressed);
370 usz_b.append_value(loc.uncompressed_size);
371 boff_b.append_value(loc.blob_offset);
372 bsz_b.append_value(loc.blob_size);
373 ck_b.append_value(loc.checksum).expect("checksum is 32 bytes");
374 }
375 let schema = lookup_schema();
376 let batch = RecordBatch::try_new(
377 schema.clone(),
378 vec![
379 Arc::new(path_b.finish()),
380 Arc::new(seq_b.finish()),
381 Arc::new(fdata_b.finish()),
382 Arc::new(comp_b.finish()),
383 Arc::new(usz_b.finish()),
384 Arc::new(boff_b.finish()),
385 Arc::new(bsz_b.finish()),
386 Arc::new(ck_b.finish()),
387 ],
388 )?;
389 self.push_subindex(&schema, &[batch], GroupKey {
390 pkg_type: RESERVED_PKG_TYPE,
391 repo: String::new(),
392 module_name: LOOKUP_MODULE.to_string(),
393 })?;
394
395 let mut builder = fst::MapBuilder::memory();
399 let mut prev: Option<&str> = None;
400 for (sorted_idx, &orig) in order.iter().enumerate() {
401 let p = self.lookup_paths[orig].as_str();
402 if prev != Some(p) {
403 builder.insert(p.as_bytes(), sorted_idx as u64)
404 .map_err(|e| anyhow!("trie insert: {e}"))?;
405 prev = Some(p);
406 }
407 }
408 let trie_bytes = builder.into_inner().map_err(|e| anyhow!("trie finish: {e}"))?;
409 self.write_raw_section(&trie_bytes, GroupKey {
410 pkg_type: RESERVED_PKG_TYPE,
411 repo: String::new(),
412 module_name: TRIE_MODULE.to_string(),
413 })
414 }
415
416 #[cfg(feature = "sign")]
421 fn write_signatures(&mut self) -> Result<()> {
422 use std::collections::BTreeMap;
423 let Some(signer) = self.signer.take() else {
424 return Ok(());
425 };
426
427 let (file_digests, artifact_paths, artifact_cms): (
430 Vec<(String, [u8; 32])>,
431 Vec<String>,
432 Vec<Vec<u8>>,
433 ) = {
434 let mut by_path: BTreeMap<&str, Vec<(u32, &[u8; 32])>> = BTreeMap::new();
435 for (p, loc) in self.lookup_paths.iter().zip(self.lookup_locs.iter()) {
436 by_path.entry(p.as_str()).or_default().push((loc.chunk_seq, &loc.checksum));
437 }
438 let mut digs = Vec::with_capacity(by_path.len());
439 let mut paths = Vec::with_capacity(by_path.len());
440 let mut cmss = Vec::with_capacity(by_path.len());
441 for (path, mut chunks) in by_path {
442 chunks.sort_by_key(|(seq, _)| *seq);
443 let n = chunks.len();
444 let digest = crate::sign::file_digest_from_parts(
445 path,
446 chunks.iter().map(|(s, c)| (*s, *c)),
447 n,
448 );
449 let cms = signer.sign_digest(&digest)?;
450 digs.push((path.to_string(), digest));
451 paths.push(path.to_string());
452 cmss.push(cms);
453 }
454 (digs, paths, cmss)
455 };
456
457 let footer = crate::index::IndexFooter::Multi { manifest_offset: 0 };
459 let root = crate::sign::archive_root(&file_digests, &footer);
460 let archive_cms = signer.sign_digest(&root)?;
461
462 let artifacts_bytes = serialize_artifact_signatures(&artifact_paths, &artifact_cms)?;
463 self.write_raw_section(
464 &artifacts_bytes,
465 GroupKey {
466 pkg_type: RESERVED_PKG_TYPE,
467 repo: String::new(),
468 module_name: SIGN_ARTIFACTS_MODULE.to_string(),
469 },
470 )?;
471 self.write_raw_section(
472 &archive_cms,
473 GroupKey {
474 pkg_type: RESERVED_PKG_TYPE,
475 repo: String::new(),
476 module_name: SIGN_ARCHIVE_MODULE.to_string(),
477 },
478 )?;
479 Ok(())
480 }
481
482 fn write_meta_subindex(&mut self) -> Result<()> {
489 let Some(table) = self.meta.take() else {
490 return Ok(());
491 };
492 let batch = build_meta_batch(&table)?;
493 let schema = meta_schema();
494 self.push_subindex(schema.as_ref(), &[batch], GroupKey {
495 pkg_type: RESERVED_PKG_TYPE,
496 repo: String::new(),
497 module_name: META_MODULE.to_string(),
498 })
499 }
500
501 fn write_raw_section(&mut self, bytes: &[u8], key: GroupKey) -> Result<()> {
504 let start = self.cursor;
505 self.file.write_all_at(bytes, start)?;
506 self.cursor += bytes.len() as u64;
507 self.entries.push(ManifestEntry {
508 pkg_type: key.pkg_type,
509 repo: key.repo,
510 module_name: key.module_name,
511 index_offset: start,
512 index_len: bytes.len() as u64,
513 row_count: 0,
514 });
515 Ok(())
516 }
517}
518
519impl ArchiveMetaSink for ArrowIpcSink {
520 fn push_subindex(
521 &mut self,
522 schema: &Schema,
523 batches: &[RecordBatch],
524 key: GroupKey,
525 ) -> Result<()> {
526 let sub_start = self.cursor;
527 let mut sub_bytes: Vec<u8> = Vec::new();
528 let mut sw = StreamWriter::try_new(&mut sub_bytes, schema)
529 .map_err(|e| anyhow!("sub-index writer: {e}"))?;
530 let mut row_count = 0u64;
531 for batch in batches {
532 row_count += batch.num_rows() as u64;
533 sw.write(batch).map_err(|e| anyhow!("sub-index write: {e}"))?;
534 }
535 sw.finish().map_err(|e| anyhow!("sub-index finish: {e}"))?;
536
537 if !is_reserved_module(&key.module_name) {
546 for batch in batches {
547 self.accumulate_lookup(batch);
548 }
549 }
550
551 let sub_len = sub_bytes.len() as u64;
552 self.file.write_all_at(&sub_bytes, sub_start)?;
553 self.cursor += sub_len;
554
555 self.entries.push(ManifestEntry {
556 pkg_type: key.pkg_type,
557 repo: key.repo,
558 module_name: key.module_name,
559 index_offset: sub_start,
560 index_len: sub_len,
561 row_count,
562 });
563 Ok(())
564 }
565
566 fn finish(mut self: Box<Self>) -> Result<u64> {
567 let order = self.lookup_order();
570 self.write_lookup_and_trie(&order)?;
571
572 self.write_reserved_sections(&order)?;
576
577 self.write_meta_subindex()?;
581
582 #[cfg(feature = "sign")]
586 self.write_signatures()?;
587
588 let manifest_offset = self.cursor;
589 let manifest_bytes =
590 write_manifest_bytes(&self.entries).map_err(|e| anyhow!("manifest: {e}"))?;
591 self.file.write_all_at(&manifest_bytes, manifest_offset)?;
592
593 let after = manifest_offset + manifest_bytes.len() as u64;
594 self.file.write_all_at(&MULTI_INDEX_MAGIC, after)?;
595 self.file.write_all_at(
596 &manifest_offset.to_le_bytes(),
597 after + MULTI_INDEX_MAGIC.len() as u64,
598 )?;
599 self.file.sync_all()?;
600
601 Ok(after + MULTI_INDEX_MAGIC.len() as u64 + 8)
602 }
603}
604
605#[cfg(feature = "sign")]
609fn serialize_artifact_signatures(paths: &[String], cms: &[Vec<u8>]) -> Result<Vec<u8>> {
610 use arrow::array::BinaryBuilder;
611 use arrow::datatypes::{DataType, Field, Schema};
612
613 let n = paths.len();
614 let schema = Arc::new(Schema::new(vec![
615 Field::new("relative_path", DataType::Utf8, false),
616 Field::new("cms", DataType::Binary, false),
617 ]));
618 let mut path_b = StringBuilder::with_capacity(n, n * 32);
619 let mut cms_b = BinaryBuilder::with_capacity(n, n * 512);
620 for (p, c) in paths.iter().zip(cms.iter()) {
621 path_b.append_value(p);
622 cms_b.append_value(c);
623 }
624 let batch = RecordBatch::try_new(
625 schema.clone(),
626 vec![Arc::new(path_b.finish()), Arc::new(cms_b.finish())],
627 )?;
628 let mut buf = Vec::new();
629 {
630 let mut w = StreamWriter::try_new(&mut buf, &schema)
631 .map_err(|e| anyhow!("artifact-sig writer: {e}"))?;
632 w.write(&batch).map_err(|e| anyhow!("artifact-sig write: {e}"))?;
633 w.finish().map_err(|e| anyhow!("artifact-sig finish: {e}"))?;
634 }
635 Ok(buf)
636}