systemprompt_loader/bundle/
extract.rs1use std::fs;
19use std::io::Read;
20use std::path::{Component, Path, PathBuf};
21
22use flate2::read::GzDecoder;
23use systemprompt_models::services::bundle::BUNDLE_MANIFEST_FILE;
24use tar::Archive;
25
26use super::error::{BundleError, BundleResult};
27
28pub const BUNDLE_TREE_PREFIX: &str = "services";
29
30#[derive(Debug, Clone, Copy, PartialEq, Eq)]
31pub enum TarLayout {
32 Bundle,
33 Root,
34}
35
36#[derive(Debug, Clone, Copy)]
37pub struct ExtractOptions<'a> {
38 pub allowed_dirs: &'a [&'a str],
39 pub max_bytes: u64,
40 pub layout: TarLayout,
41}
42
43pub fn extract_tarball(
44 archive: &Path,
45 dest: &Path,
46 opts: &ExtractOptions<'_>,
47) -> BundleResult<Vec<PathBuf>> {
48 let file = fs::File::open(archive)?;
49 extract_reader(file, dest, opts)
50}
51
52pub fn extract_bytes(
53 data: &[u8],
54 dest: &Path,
55 opts: &ExtractOptions<'_>,
56) -> BundleResult<Vec<PathBuf>> {
57 extract_reader(data, dest, opts)
58}
59
60fn extract_reader<R: Read>(
61 reader: R,
62 dest: &Path,
63 opts: &ExtractOptions<'_>,
64) -> BundleResult<Vec<PathBuf>> {
65 fs::create_dir_all(dest)?;
66 let root = dest.canonicalize()?;
67 let mut archive = Archive::new(GzDecoder::new(reader));
68 let mut written = Vec::new();
69 let mut budget = opts.max_bytes;
70
71 for entry in archive.entries()? {
72 let mut entry = entry?;
73 let entry_type = entry.header().entry_type();
74 let raw = entry.path()?.into_owned();
75
76 if !(entry_type.is_file() || entry_type.is_dir()) {
77 return Err(BundleError::extract(
78 &raw,
79 format!("disallowed entry type {entry_type:?}"),
80 ));
81 }
82 reject_traversal(&raw)?;
83
84 let size = entry.header().size().unwrap_or(0);
85 budget = budget.checked_sub(size).ok_or(BundleError::TooLarge {
86 bytes: opts.max_bytes,
87 })?;
88
89 let Some(relative) = target_path(&raw, opts)? else {
90 continue;
91 };
92 let dest_path = root.join(&relative);
93 if !dest_path.starts_with(&root) {
94 return Err(BundleError::extract(&raw, "path escapes the target"));
95 }
96
97 if entry_type.is_dir() {
98 fs::create_dir_all(&dest_path)?;
99 continue;
100 }
101 if let Some(parent) = dest_path.parent() {
102 fs::create_dir_all(parent)?;
103 }
104 entry
105 .unpack(&dest_path)
106 .map_err(|e| BundleError::extract(&dest_path, e))?;
107 written.push(relative);
108 }
109
110 Ok(written)
111}
112
113fn reject_traversal(path: &Path) -> BundleResult<()> {
114 let traversing = path.is_absolute()
115 || path
116 .components()
117 .any(|c| matches!(c, Component::ParentDir | Component::RootDir));
118 if traversing {
119 return Err(BundleError::extract(path, "invalid path in archive"));
120 }
121 Ok(())
122}
123
124fn target_path(raw: &Path, opts: &ExtractOptions<'_>) -> BundleResult<Option<PathBuf>> {
125 let normal: Vec<&str> = raw
126 .components()
127 .filter_map(|c| match c {
128 Component::Normal(s) => s.to_str(),
129 _ => None,
130 })
131 .collect();
132
133 let Some(first) = normal.first().copied() else {
134 return Ok(None);
135 };
136
137 let (stripped, rest): (PathBuf, &[&str]) = match opts.layout {
138 TarLayout::Root => (normal.iter().collect(), normal.as_slice()),
139 TarLayout::Bundle => {
140 if normal.len() == 1 && first == BUNDLE_MANIFEST_FILE {
141 return Ok(Some(PathBuf::from(BUNDLE_MANIFEST_FILE)));
142 }
143 if first != BUNDLE_TREE_PREFIX {
144 return Err(BundleError::extract(
145 raw,
146 "bundle entries must be bundle.json or under services/",
147 ));
148 }
149 (normal[1..].iter().collect(), &normal[1..])
150 },
151 };
152
153 let Some(top) = rest.first().copied() else {
154 return Ok(None);
155 };
156 if !opts.allowed_dirs.contains(&top) {
157 return Err(BundleError::extract(
158 raw,
159 format!("path not in an allowed top-level directory: {top}"),
160 ));
161 }
162 Ok(Some(stripped))
163}