1use super::S3Client;
3use anyhow::{anyhow, bail, ensure, Result};
4use aws_smithy_types::byte_stream::ByteStream;
5#[cfg(unix)]
6use std::os::unix::fs::OpenOptionsExt;
7use std::{
8 path::{Path, PathBuf},
9 sync::atomic::{AtomicU64, Ordering},
10};
11use tokio::io::AsyncWriteExt;
12
13pub const PREVIEW_LIMIT: usize = 1024 * 1024;
14
15#[derive(Debug, Clone, Copy, PartialEq, Eq)]
16pub enum TextFormat {
17 Json,
18 Yaml,
19 Plain,
20}
21
22#[derive(Debug)]
23pub struct TextPreview {
24 pub text: String,
25 pub format: TextFormat,
26 pub truncated: bool,
27}
28
29impl S3Client {
30 pub async fn preview_object(
31 &self,
32 bucket: &str,
33 region: &str,
34 key: &str,
35 ) -> Result<TextPreview> {
36 let client = self.config.s3_client_with_region(region).await;
37 let response = client
38 .get_object()
39 .bucket(bucket)
40 .key(key)
41 .range(format!("bytes=0-{PREVIEW_LIMIT}"))
42 .send()
43 .await;
44 let response = match response {
45 Ok(response) => response,
46 Err(error)
47 if error
48 .raw_response()
49 .is_some_and(|response| response.status().as_u16() == 416) =>
50 {
51 let metadata = client.head_object().bucket(bucket).key(key).send().await?;
52 ensure!(
53 metadata.content_length() == Some(0),
54 "Object range is unavailable; try downloading it"
55 );
56 return decode_preview(
57 key,
58 metadata.content_type().unwrap_or_default(),
59 Vec::new(),
60 );
61 }
62 Err(error) => return Err(error.into()),
63 };
64 let content_type = response.content_type().unwrap_or_default().to_string();
65 decode_preview(key, &content_type, read_prefix(response.body).await?)
66 }
67
68 pub async fn download_object(
69 &self,
70 bucket: &str,
71 region: &str,
72 key: &str,
73 directory: &Path,
74 ) -> Result<PathBuf> {
75 let destination = directory.join(download_name(key)?);
76 ensure!(
77 !destination.try_exists()?,
78 "File already exists: {}",
79 destination.display()
80 );
81 let client = self.config.s3_client_with_region(region).await;
82 let response = client.get_object().bucket(bucket).key(key).send().await?;
83 save_stream(response.body, response.content_length, destination).await
84 }
85}
86
87async fn read_prefix(mut body: ByteStream) -> Result<Vec<u8>> {
88 let mut bytes = Vec::new();
89 while let Some(chunk) = body.try_next().await? {
90 let remaining = (PREVIEW_LIMIT + 1).saturating_sub(bytes.len());
91 bytes.extend_from_slice(&chunk[..chunk.len().min(remaining)]);
92 if bytes.len() > PREVIEW_LIMIT {
93 break;
94 }
95 }
96 Ok(bytes)
97}
98
99fn decode_preview(key: &str, content_type: &str, mut bytes: Vec<u8>) -> Result<TextPreview> {
100 let truncated = bytes.len() > PREVIEW_LIMIT;
101 bytes.truncate(PREVIEW_LIMIT);
102 let text = match std::str::from_utf8(&bytes) {
103 Ok(text) => text,
104 Err(error) if truncated && error.error_len().is_none() => {
105 std::str::from_utf8(&bytes[..error.valid_up_to()])?
106 }
107 Err(_) => bail!("Preview supports UTF-8 text only; use Ctrl+D to download this object"),
108 };
109 ensure!(
110 !text
111 .chars()
112 .any(|c| c.is_control() && !matches!(c, '\n' | '\r' | '\t')),
113 "Binary or control-character content cannot be previewed; use Ctrl+D to download"
114 );
115 let text = text.trim_start_matches('\u{feff}');
116 let extension = key
117 .rsplit('.')
118 .next()
119 .unwrap_or_default()
120 .to_ascii_lowercase();
121 let mime = content_type
122 .split(';')
123 .next()
124 .unwrap_or_default()
125 .trim()
126 .to_ascii_lowercase();
127 ensure!(
128 mime.is_empty()
129 || mime.starts_with("text/")
130 || mime.ends_with("+json")
131 || mime.ends_with("+xml")
132 || matches!(
133 mime.as_str(),
134 "application/octet-stream"
135 | "application/json"
136 | "application/yaml"
137 | "application/x-yaml"
138 | "application/xml"
139 | "application/javascript"
140 | "application/toml"
141 | "application/x-sh"
142 ),
143 "Preview supports text formats only ({mime}); use Ctrl+D to download"
144 );
145 ensure!(
146 !matches!(
147 extension.as_str(),
148 "pdf"
149 | "zip"
150 | "gz"
151 | "tgz"
152 | "png"
153 | "jpg"
154 | "jpeg"
155 | "gif"
156 | "webp"
157 | "parquet"
158 | "avro"
159 | "exe"
160 | "wasm"
161 | "mp4"
162 | "mp3"
163 ),
164 "Preview supports text formats only; use Ctrl+D to download"
165 );
166 let format = if extension == "json" || mime == "application/json" || mime.ends_with("+json") {
167 TextFormat::Json
168 } else if matches!(extension.as_str(), "yaml" | "yml")
169 || matches!(
170 mime.as_str(),
171 "application/yaml" | "application/x-yaml" | "text/yaml"
172 )
173 {
174 TextFormat::Yaml
175 } else {
176 TextFormat::Plain
177 };
178 let formatted = if format == TextFormat::Json && !truncated {
180 serde_json::from_str::<serde_json::Value>(text)
181 .ok()
182 .and_then(|value| serde_json::to_string_pretty(&value).ok())
183 .filter(|pretty| pretty.len() <= PREVIEW_LIMIT)
184 } else {
185 None
186 };
187 Ok(TextPreview {
188 text: formatted.unwrap_or_else(|| text.into()),
189 format,
190 truncated,
191 })
192}
193
194fn download_name(key: &str) -> Result<String> {
195 let name = key.rsplit(['/', '\\']).next().unwrap_or_default();
196 ensure!(
197 !name.is_empty() && name != "." && name != "..",
198 "Object has no safe filename"
199 );
200 let name: String = name
202 .chars()
203 .map(|c| {
204 if c.is_control() || matches!(c, ':' | '*' | '?' | '"' | '<' | '>' | '|') {
205 '_'
206 } else {
207 c
208 }
209 })
210 .collect();
211 Ok(name)
212}
213
214struct PartialFile(PathBuf);
215impl Drop for PartialFile {
216 fn drop(&mut self) {
217 let _ = std::fs::remove_file(&self.0);
218 }
219}
220static DOWNLOAD_ID: AtomicU64 = AtomicU64::new(0);
221
222async fn save_stream(
223 mut body: ByteStream,
224 expected: Option<i64>,
225 destination: PathBuf,
226) -> Result<PathBuf> {
227 let parent = destination
228 .parent()
229 .ok_or_else(|| anyhow!("Download directory is missing"))?;
230 let mut opened = None;
231 for _ in 0..100 {
232 let path = parent.join(format!(
233 ".rusticity-download-{}-{}.part",
234 std::process::id(),
235 DOWNLOAD_ID.fetch_add(1, Ordering::Relaxed)
236 ));
237 let mut options = std::fs::OpenOptions::new();
238 options.write(true).create_new(true);
239 #[cfg(unix)]
240 options.mode(0o600);
241 match options.open(&path) {
243 Ok(file) => {
244 opened = Some((PartialFile(path), tokio::fs::File::from_std(file)));
245 break;
246 }
247 Err(error) if error.kind() == std::io::ErrorKind::AlreadyExists => continue,
248 Err(error) => return Err(error.into()),
249 }
250 }
251 let (partial, mut file) =
252 opened.ok_or_else(|| anyhow!("Could not create a unique partial download"))?;
253 let mut written = 0u64;
254 while let Some(chunk) = body.try_next().await? {
255 file.write_all(&chunk).await?;
256 written += chunk.len() as u64;
257 }
258 if let Some(expected) = expected {
259 ensure!(
260 expected >= 0 && written == expected as u64,
261 "Incomplete download: received {written} of {expected} bytes"
262 );
263 }
264 file.flush().await?;
265 file.sync_all().await?;
266 drop(file);
267 std::fs::hard_link(&partial.0, &destination).map_err(|error| {
272 anyhow!(
273 "Cannot save {} (existing files are never overwritten): {error}",
274 destination.display()
275 )
276 })?;
277 Ok(destination)
278}
279
280#[cfg(test)]
281mod tests {
282 use super::*;
283 fn directory() -> PathBuf {
284 let path = std::env::temp_dir().join(format!(
285 "rusticity-download-test-{}-{}",
286 std::process::id(),
287 DOWNLOAD_ID.fetch_add(1, Ordering::Relaxed)
288 ));
289 std::fs::create_dir(&path).unwrap();
290 path
291 }
292 #[test]
293 fn text_formats_binary_and_utf8_boundaries() {
294 let json = decode_preview("dir/object.json", "", br#"{"hello":true}"#.to_vec()).unwrap();
295 assert_eq!(json.format, TextFormat::Json);
296 assert!(json.text.contains("\n \"hello\": true"));
297 assert_eq!(
298 decode_preview("x.YML", "", b"enabled: true".to_vec())
299 .unwrap()
300 .format,
301 TextFormat::Yaml
302 );
303 assert_eq!(
304 decode_preview("x.txt", "text/plain", b"hello".to_vec())
305 .unwrap()
306 .format,
307 TextFormat::Plain
308 );
309 assert!(decode_preview("x.json", "", vec![0xff]).is_err());
310 assert!(decode_preview("x.txt", "", b"bad\x1b[2J".to_vec()).is_err());
311 assert!(decode_preview("x.bin", "", b"abc\0def".to_vec()).is_err());
312 assert_eq!(decode_preview("empty.txt", "", vec![]).unwrap().text, "");
313 assert!(decode_preview("x.txt", "application/pdf", b"%PDF-ascii".to_vec()).is_err());
314 assert!(
315 decode_preview("x.pdf", "application/octet-stream", b"%PDF-ascii".to_vec()).is_err()
316 );
317 let mut bytes = vec![b'a'; PREVIEW_LIMIT - 1];
318 bytes.extend_from_slice("界".as_bytes());
319 let preview = decode_preview("x.txt", "", bytes).unwrap();
320 assert!(preview.truncated);
321 assert_eq!(preview.text.len(), PREVIEW_LIMIT - 1);
322 }
323 #[test]
324 fn object_keys_cannot_escape_download_directory() {
325 assert_eq!(
326 download_name("../../some/object.json").unwrap(),
327 "object.json"
328 );
329 assert_eq!(download_name("a\\b.txt").unwrap(), "b.txt");
330 for key in ["", "folder/", ".", "..", "folder/.."] {
331 assert!(download_name(key).is_err());
332 }
333 assert_eq!(download_name("bad\nname.txt").unwrap(), "bad_name.txt");
334 }
335 #[tokio::test]
336 async fn download_publishes_exact_bytes_and_never_overwrites() {
337 let directory = directory();
338 let target = directory.join("object.txt");
339 assert_eq!(
340 save_stream(ByteStream::from_static(b"hello"), Some(5), target.clone())
341 .await
342 .unwrap(),
343 target
344 );
345 assert_eq!(std::fs::read(&target).unwrap(), b"hello");
346 assert!(
347 save_stream(ByteStream::from_static(b"different"), None, target.clone())
348 .await
349 .is_err()
350 );
351 assert_eq!(std::fs::read(&target).unwrap(), b"hello");
352 assert_eq!(std::fs::read_dir(&directory).unwrap().count(), 1);
353 std::fs::remove_dir_all(directory).unwrap();
354 }
355 #[tokio::test]
356 async fn incomplete_download_removes_partial_and_does_not_publish() {
357 let directory = directory();
358 let target = directory.join("object.txt");
359 assert!(save_stream(
360 ByteStream::from_static(b"partial"),
361 Some(100),
362 target.clone()
363 )
364 .await
365 .is_err());
366 assert!(!target.exists());
367 assert_eq!(std::fs::read_dir(&directory).unwrap().count(), 0);
368 std::fs::remove_dir_all(directory).unwrap();
369 }
370 #[tokio::test]
371 async fn cancelled_download_removes_its_partial_file() {
372 let directory = directory();
373 let target = directory.join("large.bin");
374 let destination = target.clone();
375 let task = tokio::spawn(async move {
376 save_stream(
377 ByteStream::from(vec![42; 8 * 1024 * 1024]),
378 None,
379 destination,
380 )
381 .await
382 });
383 tokio::time::timeout(std::time::Duration::from_secs(5), async {
384 while std::fs::read_dir(&directory).unwrap().count() == 0 {
385 tokio::task::yield_now().await;
386 }
387 })
388 .await
389 .unwrap();
390 task.abort();
391 assert!(task.await.unwrap_err().is_cancelled());
392 assert!(!target.exists());
393 assert_eq!(std::fs::read_dir(&directory).unwrap().count(), 0);
394 std::fs::remove_dir_all(directory).unwrap();
395 }
396
397 #[cfg(unix)]
398 #[tokio::test]
399 async fn downloads_do_not_replace_symlinks_and_publish_private_permissions() {
400 use std::os::unix::fs::{symlink, PermissionsExt};
401 let directory = directory();
402 let target = directory.join("object.bin");
403 let unrelated = directory.join("unrelated.bin");
404 std::fs::write(&unrelated, b"keep").unwrap();
405 symlink(&unrelated, &target).unwrap();
406 assert!(
407 save_stream(ByteStream::from_static(b"new"), None, target.clone())
408 .await
409 .is_err()
410 );
411 assert!(std::fs::symlink_metadata(&target).unwrap().is_symlink());
412 assert_eq!(std::fs::read(&unrelated).unwrap(), b"keep");
413 let private = directory.join("private.bin");
414 save_stream(ByteStream::from_static(b"new"), None, private.clone())
415 .await
416 .unwrap();
417 assert_eq!(
418 std::fs::metadata(private).unwrap().permissions().mode() & 0o777,
419 0o600
420 );
421 std::fs::remove_dir_all(directory).unwrap();
422 }
423
424 #[tokio::test]
425 async fn preview_reader_stops_at_bounded_prefix() {
426 let bytes = read_prefix(ByteStream::from(vec![b'x'; PREVIEW_LIMIT * 2]))
427 .await
428 .unwrap();
429 assert_eq!(bytes.len(), PREVIEW_LIMIT + 1);
430 }
431}