Skip to main content

rusticity_core/s3/
object.rs

1//! Bounded object previews and streaming, no-overwrite downloads.
2use 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    // Pretty-print only complete, valid JSON; preserve truncated/invalid text for inspection.
179    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    // Do not interpret S3 keys as local paths, or allow terminal controls in notices.
201    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        // Create and register cleanup without an await/cancellation gap.
242        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    // Linking in the same directory atomically publishes the completed file and
268    // fails if *anything* already occupies the destination (including a symlink).
269    // This short metadata operation is synchronous so task cancellation cannot
270    // race publication after the UI has reported cancellation.
271    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}