1use std::fs;
2use std::io::{self, Read};
3use std::path::{Component, Path, PathBuf};
4
5use tar::{Archive, EntryType};
6use tracing::warn;
7
8use crate::docs_error::DocsError;
9
10const MAX_BYTES_PER_DOCS_MB: u64 = 20 * 1_000_000;
12const MAX_ENTRIES_PER_DOCS_MB: usize = 1_000;
14
15#[derive(Debug, Clone, Copy, PartialEq, Eq)]
17pub struct ExtractLimits {
18 max_bytes: u64,
19 max_entries: usize,
20}
21
22impl ExtractLimits {
23 pub fn new(max_bytes: u64, max_entries: usize) -> Self {
24 Self {
25 max_bytes,
26 max_entries,
27 }
28 }
29
30 pub fn from_max_docs_size(max_size_mb: usize) -> Self {
33 let max_size_mb = max_size_mb.max(1);
34 Self::new(
35 MAX_BYTES_PER_DOCS_MB.saturating_mul(max_size_mb as u64),
36 MAX_ENTRIES_PER_DOCS_MB.saturating_mul(max_size_mb),
37 )
38 }
39
40 pub(crate) fn budget(self) -> ExtractBudget {
41 ExtractBudget {
42 remaining_bytes: self.max_bytes,
43 remaining_entries: self.max_entries,
44 }
45 }
46}
47
48pub(crate) struct ExtractBudget {
50 remaining_bytes: u64,
51 remaining_entries: usize,
52}
53
54impl ExtractBudget {
55 pub(crate) fn take_entry(&mut self) -> Result<(), DocsError> {
57 self.remaining_entries = self
58 .remaining_entries
59 .checked_sub(1)
60 .ok_or(DocsError::ArchiveTooLarge)?;
61 Ok(())
62 }
63
64 pub(crate) fn write_file(
67 &mut self,
68 content: &mut impl Read,
69 target: &Path,
70 ) -> Result<(), DocsError> {
71 if let Some(parent) = target.parent() {
72 fs::create_dir_all(parent)?;
73 }
74 let mut file = fs::File::create(target)?;
75 let written = io::copy(
76 &mut content.take(self.remaining_bytes.saturating_add(1)),
77 &mut file,
78 )?;
79 self.remaining_bytes = self
80 .remaining_bytes
81 .checked_sub(written)
82 .ok_or(DocsError::ArchiveTooLarge)?;
83 Ok(())
84 }
85}
86
87pub fn unpack_crate(
94 reader: impl Read,
95 dest: &Path,
96 root: &str,
97 limits: ExtractLimits,
98) -> Result<(), DocsError> {
99 let mut budget = limits.budget();
100 let mut archive = Archive::new(reader);
101 for entry in archive.entries()? {
102 let mut entry = entry?;
103 budget.take_entry()?;
104 let path = entry.path()?.into_owned();
105 let Some(relative) = path_below_root(&path, root) else {
106 warn!(
107 "Skipping crate archive entry outside of {root}: {}",
108 path.display()
109 );
110 continue;
111 };
112 let target = dest.join(relative);
113 match entry.header().entry_type() {
114 EntryType::Directory => fs::create_dir_all(&target)?,
115 EntryType::Regular | EntryType::Continuous => {
116 budget.write_file(&mut entry, &target)?;
117 }
118 other => warn!(
119 "Skipping crate archive entry {} of type {other:?}",
120 path.display()
121 ),
122 }
123 }
124 Ok(())
125}
126
127fn path_below_root(path: &Path, root: &str) -> Option<PathBuf> {
129 let mut components = path.components();
130 if components.next() != Some(Component::Normal(root.as_ref())) {
131 return None;
132 }
133 components
134 .all(|c| matches!(c, Component::Normal(_)))
135 .then(|| path.to_path_buf())
136}
137
138#[cfg(test)]
139mod tests {
140 use tar::{Builder, Header};
141
142 use super::*;
143
144 fn tar_header(path: &str, entry_type: EntryType, size: u64) -> Header {
145 let mut header = Header::new_gnu();
146 header.as_gnu_mut().unwrap().name[..path.len()].copy_from_slice(path.as_bytes());
148 header.set_entry_type(entry_type);
149 header.set_size(size);
150 header.set_mode(0o644);
151 header.set_cksum();
152 header
153 }
154
155 fn append_file(builder: &mut Builder<Vec<u8>>, path: &str, content: &[u8]) {
156 let header = tar_header(path, EntryType::Regular, content.len() as u64);
157 builder.append(&header, content).unwrap();
158 }
159
160 fn append_link(builder: &mut Builder<Vec<u8>>, path: &str, kind: EntryType, target: &str) {
161 let mut header = tar_header(path, kind, 0);
162 header.set_link_name(target).unwrap();
163 header.set_cksum();
164 builder.append(&header, io::empty()).unwrap();
165 }
166
167 const LIMITS: ExtractLimits = ExtractLimits {
168 max_bytes: 1_000,
169 max_entries: 10,
170 };
171
172 fn try_unpack(
173 limits: ExtractLimits,
174 build: impl FnOnce(&mut Builder<Vec<u8>>),
175 ) -> (tempfile::TempDir, Result<(), DocsError>) {
176 let mut builder = Builder::new(Vec::new());
177 build(&mut builder);
178 let tar = builder.into_inner().unwrap();
179 let dest = tempfile::tempdir().unwrap();
180 let result = unpack_crate(tar.as_slice(), dest.path(), "x-1.0.0", limits);
181 (dest, result)
182 }
183
184 fn unpack(build: impl FnOnce(&mut Builder<Vec<u8>>)) -> tempfile::TempDir {
185 let (dest, result) = try_unpack(LIMITS, build);
186 result.unwrap();
187 dest
188 }
189
190 #[test]
191 fn unpack_crate_fails_when_uncompressed_size_exceeds_budget() {
192 let (_dest, result) = try_unpack(LIMITS, |b| {
193 append_file(b, "x-1.0.0/a", &[0; 600]);
194 append_file(b, "x-1.0.0/b", &[0; 600]);
195 });
196
197 assert!(matches!(result, Err(DocsError::ArchiveTooLarge)));
198 }
199
200 #[test]
201 fn unpack_crate_accepts_archive_of_exactly_the_budget() {
202 let (_dest, result) = try_unpack(LIMITS, |b| {
203 append_file(b, "x-1.0.0/a", &[0; 500]);
204 append_file(b, "x-1.0.0/b", &[0; 500]);
205 });
206
207 assert!(result.is_ok());
208 }
209
210 #[test]
211 fn unpack_crate_fails_when_entry_count_exceeds_budget() {
212 let (_dest, result) = try_unpack(LIMITS, |b| {
213 for i in 0..11 {
214 append_file(b, &format!("x-1.0.0/{i}"), b"");
215 }
216 });
217
218 assert!(matches!(result, Err(DocsError::ArchiveTooLarge)));
219 }
220
221 #[test]
222 fn extract_limits_scale_with_max_docs_size() {
223 assert_eq!(
224 ExtractLimits::from_max_docs_size(100),
225 ExtractLimits::new(2_000_000_000, 100_000)
226 );
227 assert_eq!(
228 ExtractLimits::from_max_docs_size(0),
229 ExtractLimits::from_max_docs_size(1)
230 );
231 }
232
233 #[test]
234 fn unpack_crate_writes_files_and_directories() {
235 let dest = unpack(|b| {
236 let dir = tar_header("x-1.0.0/src", EntryType::Directory, 0);
237 b.append(&dir, io::empty()).unwrap();
238 append_file(b, "x-1.0.0/Cargo.toml", b"[package]");
239 append_file(b, "x-1.0.0/src/lib.rs", b"pub fn f() {}");
240 });
241
242 let root = dest.path().join("x-1.0.0");
243 assert_eq!(fs::read(root.join("Cargo.toml")).unwrap(), b"[package]");
244 assert_eq!(fs::read(root.join("src/lib.rs")).unwrap(), b"pub fn f() {}");
245 }
246
247 #[test]
248 fn unpack_crate_skips_symlinks_and_hardlinks() {
249 let secret = tempfile::NamedTempFile::new().unwrap();
250 fs::write(secret.path(), b"secret").unwrap();
251 let target = secret.path().to_str().unwrap();
252
253 let dest = unpack(|b| {
254 append_link(b, "x-1.0.0/symlink.html", EntryType::Symlink, target);
255 append_link(b, "x-1.0.0/hardlink.html", EntryType::Link, target);
256 append_file(b, "x-1.0.0/Cargo.toml", b"[package]");
257 });
258
259 let root = dest.path().join("x-1.0.0");
260 assert!(fs::symlink_metadata(root.join("symlink.html")).is_err());
261 assert!(fs::symlink_metadata(root.join("hardlink.html")).is_err());
262 assert!(root.join("Cargo.toml").exists());
263 }
264
265 #[test]
266 fn unpack_crate_skips_entries_outside_of_root() {
267 let dest = unpack(|b| {
268 append_file(b, "other/file", b"other");
269 append_file(b, "x-1.0.0/../escape", b"escape");
270 append_file(b, "/abs", b"abs");
271 append_file(b, "x-1.0.0/Cargo.toml", b"[package]");
272 });
273
274 assert!(!dest.path().join("other").exists());
275 assert!(!dest.path().join("escape").exists());
276 assert!(dest.path().join("x-1.0.0/Cargo.toml").exists());
277 let entries: Vec<_> = fs::read_dir(dest.path()).unwrap().collect();
278 assert_eq!(entries.len(), 1);
279 }
280}